vllm-omni/tests/worker/test_native_duplex_hooks.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

import ast
import inspect
from dataclasses import replace
from types import SimpleNamespace

import numpy as np
import pytest
import torch

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


def _minicpmo_duplex_policy_case(
    state: SimpleNamespace,
    payload: dict[str, object],
):
    from vllm_omni.model_executor.duplex_sampling import DuplexSamplingRow
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    session_key = "sid-policy"
    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model._minicpmo45_native_duplex_token_ids_cache = {
        "listen_token_id": 7,
        "tts_bos_token_id": 8,
        "turn_eos_token_id": 9,
    }
    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state})
    row = DuplexSamplingRow(
        row_idx=0,
        request_id="req-policy",
        session_id=session_key,
        seq=3,
        payload=payload,
        max_tokens=20,
    )
    return model, row


def test_minicpmo_model_hook_owns_duplex_sampling_rows_and_force_listen():
    from vllm_omni.model_executor.duplex_sampling import DuplexSamplingRow
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    listen_id = 7
    state = SimpleNamespace(
        current_turn_ended=True,
        last_terminator_token=None,
        pending_terminator_token=None,
    )
    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model._minicpmo45_native_duplex_token_ids_cache = {
        "listen_token_id": listen_id,
        "tts_bos_token_id": 8,
        "turn_eos_token_id": 9,
    }
    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={"sid-hook": state})
    logits = torch.zeros((1, 16), dtype=torch.float32)
    row = DuplexSamplingRow(
        row_idx=0,
        request_id="req-hook",
        session_id="sid-hook",
        seq=3,
        payload={"force_listen": True, "is_speech": True},
        max_tokens=20,
    )

    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))

    assert model._minicpmo45_active_duplex_rows == [0]
    assert model._minicpmo45_duplex_row_sessions == {0: "sid-hook"}
    assert model._minicpmo45_duplex_row_payloads == {0: row.payload}
    assert model._minicpmo45_duplex_row_max_tokens == {0: 20}
    assert logits[0, listen_id].item() == 0.0
    assert torch.isneginf(logits[0, :listen_id]).all()
    assert torch.isneginf(logits[0, listen_id + 1 :]).all()


def test_minicpmo_model_hook_turn_end_latch_forces_silence_to_listen():
    state = SimpleNamespace(
        current_turn_ended=True,
        last_terminator_token=9,
        pending_terminator_token=None,
    )
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False})
    logits = torch.zeros((1, 16), dtype=torch.float32)
    logits[0, 10] = 20.0

    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))

    assert logits[0, 7].item() == 0.0
    assert torch.isneginf(logits[0, :7]).all()
    assert torch.isneginf(logits[0, 8:]).all()


def test_minicpmo_model_hook_pending_speech_after_turn_eos_allows_silence_sampling():
    state = SimpleNamespace(
        current_turn_ended=True,
        last_terminator_token=9,
        pending_terminator_token=None,
        pending_speech_context=True,
    )
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False})
    logits = torch.zeros((1, 16), dtype=torch.float32)
    logits[0, 7] = 30.0
    logits[0, 10] = 20.0
    original_logits = logits.clone()

    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))

    assert torch.equal(logits, original_logits)
    assert state.pending_speech_context is True


def test_minicpmo_model_hook_speech_row_does_not_create_pending_context():
    state = SimpleNamespace(
        current_turn_ended=True,
        last_terminator_token=9,
        pending_terminator_token=None,
        pending_speech_context=False,
    )
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": True})
    logits = torch.zeros((1, 16), dtype=torch.float32)
    logits[0, 7] = 30.0
    logits[0, 8] = -2.0
    logits[0, 10] = 20.0
    original_logits = logits.clone()

    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))

    assert state.current_turn_ended is True
    assert state.last_terminator_token is None
    assert state.pending_speech_context is False
    assert torch.equal(logits, original_logits)


def test_minicpmo_model_hook_old_response_output_does_not_clear_pending_speech_context():
    state = SimpleNamespace(
        current_turn_ended=False,
        last_terminator_token=None,
        pending_terminator_token=None,
        pending_speech_context=True,
    )
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": True})
    model._minicpmo45_duplex_row_sessions = {0: row.session_id}
    model._minicpmo45_duplex_row_payloads = {0: row.payload}

    model._record_minicpmo45_duplex_terminator(
        0,
        10,
        {"listen_token_id": 7, "chunk_eos_token_id": -1, "chunk_tts_eos_token_id": -1, "turn_eos_token_id": 9},
    )

    assert state.pending_speech_context is True
    assert state.current_turn_ended is False


def test_minicpmo_model_hook_new_response_output_clears_pending_speech_context():
    state = SimpleNamespace(
        current_turn_ended=True,
        last_terminator_token=7,
        pending_terminator_token=7,
        pending_speech_context=True,
    )
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False})
    model._minicpmo45_duplex_row_sessions = {0: row.session_id}
    model._minicpmo45_duplex_row_payloads = {0: row.payload}

    model._record_minicpmo45_duplex_terminator(
        0,
        8,
        {"listen_token_id": 7, "chunk_eos_token_id": -1, "chunk_tts_eos_token_id": -1, "turn_eos_token_id": 9},
    )

    assert state.pending_speech_context is False
    assert state.current_turn_ended is False


def test_minicpmo_model_hook_empty_speak_envelope_preserves_pending_speech_context():
    state = SimpleNamespace(
        current_turn_ended=True,
        last_terminator_token=9,
        pending_terminator_token=None,
        pending_speech_context=True,
        pending_speech_response_open=False,
    )
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False})
    model._minicpmo45_duplex_row_sessions = {0: row.session_id}
    model._minicpmo45_duplex_row_payloads = {0: row.payload}
    token_ids = {
        "listen_token_id": 7,
        "tts_bos_token_id": 8,
        "chunk_eos_token_id": -1,
        "chunk_tts_eos_token_id": -1,
        "turn_eos_token_id": 9,
    }

    model._record_minicpmo45_duplex_terminator(0, 8, token_ids)
    model._record_minicpmo45_duplex_terminator(0, 9, token_ids)

    assert state.current_turn_ended is True
    assert state.pending_speech_context is True
    assert state.pending_speech_response_open is False

    logits = torch.zeros((1, 16), dtype=torch.float32)
    logits[0, 7] = 30.0
    logits[0, 10] = 20.0
    original_logits = logits.clone()

    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))

    assert torch.equal(logits, original_logits)


def test_minicpmo_model_hook_second_new_response_step_is_not_forced_to_listen():
    state = SimpleNamespace(
        current_turn_ended=True,
        last_terminator_token=9,
        pending_terminator_token=None,
        pending_speech_context=True,
    )
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False})
    logits = torch.zeros((1, 16), dtype=torch.float32)
    logits[0, 7] = 30.0
    logits[0, 10] = 20.0
    original_logits = logits.clone()

    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
    model._minicpmo45_duplex_row_sessions = {0: row.session_id}
    model._minicpmo45_duplex_row_payloads = {0: row.payload}
    model._record_minicpmo45_duplex_terminator(
        0,
        8,
        {"listen_token_id": 7, "chunk_eos_token_id": -1, "chunk_tts_eos_token_id": -1, "turn_eos_token_id": 9},
    )
    next_logits = torch.zeros((1, 16), dtype=torch.float32)
    next_logits[0, 7] = 30.0
    next_logits[0, 10] = 20.0
    next_original_logits = next_logits.clone()

    model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (row,))

    assert torch.equal(logits, original_logits)
    assert torch.equal(next_logits, next_original_logits)
    assert state.pending_speech_context is False
    assert state.current_turn_ended is False


def test_minicpmo_model_hook_same_speech_row_second_step_does_not_rearm_pending_context():
    state = SimpleNamespace(
        current_turn_ended=True,
        last_terminator_token=9,
        pending_terminator_token=None,
        pending_speech_context=True,
    )
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": True})
    logits = torch.zeros((1, 16), dtype=torch.float32)
    logits[0, 7] = 30.0
    logits[0, 10] = 20.0

    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
    model._minicpmo45_duplex_row_sessions = {0: row.session_id}
    model._minicpmo45_duplex_row_payloads = {0: row.payload}
    model._record_minicpmo45_duplex_terminator(
        0,
        10,
        {"listen_token_id": 7, "chunk_eos_token_id": -1, "chunk_tts_eos_token_id": -1, "turn_eos_token_id": 9},
    )
    next_logits = torch.zeros((1, 16), dtype=torch.float32)
    next_logits[0, 7] = 30.0
    next_logits[0, 8] = -2.0
    next_logits[0, 10] = 20.0
    original_next_logits = next_logits.clone()

    model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (row,))

    assert state.pending_speech_context is False
    assert state.current_turn_ended is False
    assert torch.equal(next_logits, original_next_logits)


def test_minicpmo_model_hook_mid_turn_speech_preserves_logits_for_post_sample_redirect():
    state = SimpleNamespace(
        current_turn_ended=False,
        last_terminator_token=None,
        pending_terminator_token=None,
    )
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": True})
    logits = torch.zeros((1, 16), dtype=torch.float32)
    logits[0, 7] = 30.0
    logits[0, 8] = -2.0
    logits[0, 10] = 20.0
    original_logits = logits.clone()

    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))

    assert torch.equal(logits, original_logits)


def test_minicpmo_model_hook_ignores_serving_new_user_turn_marker():
    state = SimpleNamespace(
        current_turn_ended=False,
        last_terminator_token=8,
        pending_terminator_token=8,
    )
    model, row = _minicpmo_duplex_policy_case(
        state,
        {"is_speech": True, "new_user_turn": True, "force_speak": True},
    )
    logits = torch.zeros((1, 16), dtype=torch.float32)
    logits[0, 7] = 5.0
    logits[0, 10] = 20.0
    original_logits = logits.clone()
    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))

    assert state.current_turn_ended is False
    assert state.last_terminator_token is None
    assert torch.equal(logits, original_logits)


def test_generic_ar_runner_builds_typed_duplex_sampling_rows():
    from vllm_omni.model_executor.duplex_sampling import DuplexSamplingHelper
    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner

    runner = GPUARModelRunner.__new__(GPUARModelRunner)
    runner.input_batch = SimpleNamespace(req_ids=["req-duplex", "req-plain"])
    runner.model_intermediate_buffer = {
        "req-duplex": {
            "duplex": {
                "data_plane": True,
                "session_id": "sid-runner-hook",
                "seq": 4,
                "payload": {"is_speech": True},
            }
        }
    }
    runner.requests = {
        "req-duplex": SimpleNamespace(
            sampling_params=SimpleNamespace(max_tokens=32),
        )
    }
    helper = DuplexSamplingHelper()
    helper.active_request_ids = {"req-duplex"}

    rows = helper.rows(runner)

    assert len(rows) == 1
    assert rows[0].row_idx == 0
    assert rows[0].request_id == "req-duplex"
    assert rows[0].session_id == "sid-runner-hook"
    assert rows[0].seq == 4
    assert rows[0].payload == {"is_speech": True}
    assert rows[0].max_tokens == 32


def test_generic_ar_runner_skips_duplex_rows_without_model_hook():
    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner

    runner = GPUARModelRunner.__new__(GPUARModelRunner)
    runner.input_batch = SimpleNamespace(
        sampling_metadata=object(),
        update_async_output_token_ids=lambda: None,
    )
    runner.model = SimpleNamespace(sample=lambda *_args, **_kwargs: None, prefer_model_sampler=False)
    runner.sampler = lambda **_kwargs: "standard-sampler"
    runner._duplex_sampling_helper = SimpleNamespace(
        active_request_ids={"req-duplex"},
        rows=lambda *_args: pytest.fail("duplex rows built without a model hook"),
    )

    assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "standard-sampler"


def test_plain_model_hook_resolution_does_not_allocate_duplex_tracking():
    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner

    runner = GPUARModelRunner.__new__(GPUARModelRunner)
    runner.model = SimpleNamespace()
    runner._duplex_sampling_hook = None
    runner._duplex_sampling_hook_resolved = False

    assert runner._resolve_duplex_sampling_hook() is None
    assert not hasattr(runner, "_duplex_sampling_helper")


def test_minicpmo_non_duplex_sample_skips_duplex_row_scan():
    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner

    calls = []
    runner = GPUARModelRunner.__new__(GPUARModelRunner)
    runner.input_batch = SimpleNamespace(
        sampling_metadata=SimpleNamespace(output_token_ids=[]),
        update_async_output_token_ids=lambda: None,
    )
    runner.model = SimpleNamespace(
        sample=lambda *_args, **_kwargs: "model-sampler",
        prefer_model_sampler=True,
        skips_model_sampler_output_token_history=True,
        prepare_duplex_sampling=lambda *_args: calls.append("prepare"),
    )
    runner.sampler = SimpleNamespace()
    runner._duplex_sampling_helper = SimpleNamespace(
        active_request_ids=set(),
        hook_active=False,
        rows=lambda *_args: pytest.fail("non-duplex MiniCPM scanned request rows"),
    )

    assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler"
    assert calls == []


def test_minicpmo_duplex_sample_clears_stale_rows_once_without_scanning():
    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner

    calls = []
    runner = GPUARModelRunner.__new__(GPUARModelRunner)
    runner.input_batch = SimpleNamespace(
        sampling_metadata=SimpleNamespace(output_token_ids=[]),
        update_async_output_token_ids=lambda: None,
    )
    runner.model = SimpleNamespace(
        sample=lambda *_args, **_kwargs: "model-sampler",
        prefer_model_sampler=True,
        skips_model_sampler_output_token_history=True,
        prepare_duplex_sampling=lambda _logits, _metadata, rows: calls.append(rows),
    )
    runner.sampler = SimpleNamespace()
    runner._duplex_sampling_helper = SimpleNamespace(
        active_request_ids=set(),
        hook_active=True,
        rows=lambda *_args: pytest.fail("duplex cleanup scanned request rows"),
    )

    assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler"
    assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler"
    assert calls == [()]


def test_generic_ar_runner_has_no_minicpmo_sampler_state_or_typeerror_probe():
    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner

    source = inspect.getsource(GPUARModelRunner)

    assert "_minicpmo45_duplex_row" not in source
    assert "_minicpmo45_native_duplex_token_ids" not in source
    assert 'if "duplex_rows" not in str(exc)' not in source


def test_minicpmo_model_cleans_session_state_when_request_finishes():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    request_id = "duplex-sid-cleanup-i3-e0-stage0"
    session_key = "sid-cleanup"
    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model = SimpleNamespace()
    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: object()})
    model._minicpmo45_duplex_request_sessions = {request_id: session_key}
    model._minicpmo45_force_listen_applied_segments = {
        (request_id, 1),
        ("duplex-sid-other-e0-stage0", 2),
    }

    model.on_requests_finished({request_id})

    assert model._minicpmo45_duplex_data_plane_helper.sessions == {}
    assert model._minicpmo45_duplex_request_sessions == {}
    assert model._minicpmo45_force_listen_applied_segments == {
        ("duplex-sid-other-e0-stage0", 2),
    }


def test_minicpmo_stage0_routes_duplex_metadata_per_batched_request():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )
    from vllm_omni.utils.mm_outputs import to_payload_element

    class _Thinker(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.register_buffer("device_anchor", torch.zeros(1))

        def forward(self, *, input_ids, **kwargs):
            del kwargs
            return torch.arange(input_ids.numel() * 4, dtype=torch.float32).reshape(input_ids.numel(), 4)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    torch.nn.Module.__init__(model)
    model.model_stage = "llm"
    model.thinker = _Thinker()
    output = model.forward(
        input_ids=torch.tensor([1, 2, 3, 4]),
        positions=torch.arange(4),
        runtime_additional_information=[
            {
                "global_request_id": ["req-a"],
                "duplex": {
                    "duplex_prompt_token_ids": [101, 102],
                    "special_token_ids": {"listen_token_id": 701},
                },
            },
            {
                "global_request_id": ["req-b"],
                "duplex": {
                    "duplex_prompt_token_ids": [201, 202, 203],
                    "special_token_ids": {"listen_token_id": 702},
                },
            },
        ],
    )

    prompt_rows = output.multimodal_outputs["duplex_prompt_token_ids"]
    listen_rows = output.multimodal_outputs["meta"]["listen_token_id"]
    assert to_payload_element(prompt_rows, 0, 0, 2) == [101, 102]
    assert to_payload_element(prompt_rows, 1, 2, 4) == [201, 202, 203]
    assert int(to_payload_element(listen_rows, 0, 0, 2).reshape(-1)[0]) == 701
    assert int(to_payload_element(listen_rows, 1, 2, 4).reshape(-1)[0]) == 702


def test_minicpmo_stage0_rejects_invalid_resolved_ref_audio():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
    )

    with pytest.raises(ValueError, match="invalid native duplex ref_audio_data"):
        MiniCPMO45Stage0DuplexRuntime._decode_ref_audio_from_session_config(
            {
                "ref_audio_data": "a",
                "ref_audio_format": "pcm_f32le",
            }
        )


def test_minicpmo_stage0_special_token_ids_are_tokenizer_derived():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
    )

    class _Tokenizer:
        unk_token_id = 0
        ids = {
            "<unit>": 101,
            "</unit>": 102,
            "<|listen|>": 103,
            "<|speak|>": 104,
            "<|tts_bos|>": 105,
            "<|tts_eos|>": 106,
            "<|tts_pad|>": 107,
            "<|chunk_eos|>": 108,
            "<|chunk_tts_eos|>": 109,
            "<|turn_eos|>": 110,
        }

        def convert_tokens_to_ids(self, token):
            return self.ids.get(token, self.unk_token_id)

        def encode(self, text, add_special_tokens=False):
            del add_special_tokens
            return [self.ids[text]] if text in self.ids else [201, self.ids["<|tts_bos|>"]]

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.tokenizer = _Tokenizer()
    runtime._init_token_ids()

    runtime._require_special_token_ids()
    assert runtime.tts_bos_token_id == 105
    assert runtime.stage_padding_token_id() == 102
    assert runtime._special_token_ids()["chunk_tts_eos_token_id"] == 109


def test_minicpmo_stage0_rejects_unknown_special_token_fallbacks():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
    )

    class _Tokenizer:
        unk_token_id = 0
        ids = {
            "<unit>": 101,
            "</unit>": 102,
            "<|listen|>": 103,
            "<|speak|>": 104,
            "<|tts_eos|>": 106,
            "<|tts_pad|>": 107,
            "<|chunk_eos|>": 108,
            "<|chunk_tts_eos|>": 109,
            "<|turn_eos|>": 110,
        }

        def convert_tokens_to_ids(self, token):
            return self.ids.get(token, self.unk_token_id)

        def encode(self, text, add_special_tokens=False):
            del text, add_special_tokens
            return [self.unk_token_id]

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.tokenizer = _Tokenizer()
    runtime._init_token_ids()

    with pytest.raises(ValueError, match=r"<\|tts_bos\|>"):
        runtime._require_special_token_ids()


def _minicpmo45_tokenizer_stub():
    token_ids = {
        token: index
        for index, token in enumerate(
            (
                "<unit>",
                "</unit>",
                "<|listen|>",
                "<|speak|>",
                "<|tts_bos|>",
                "<|tts_eos|>",
                "<|tts_pad|>",
                "<|chunk_eos|>",
                "<|chunk_tts_eos|>",
                "<|turn_eos|>",
            ),
            start=101,
        )
    }
    return SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: token_ids.get(token, 0),
        encode=lambda token, add_special_tokens=False: (
            [token_ids[token]] if not add_special_tokens and token in token_ids else [0]
        ),
    )


def _transformers_processor_stub(auto_processor):
    class _AutoImageProcessor:
        @staticmethod
        def register(*_args, **_kwargs):
            return None

    return SimpleNamespace(
        AutoImageProcessor=_AutoImageProcessor,
        AutoProcessor=auto_processor,
    )


def test_minicpmo_stage0_loads_processor_from_hf_id(monkeypatch):
    import sys

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
    )

    model_id = "openbmb/MiniCPM-o-4_5"
    processor = SimpleNamespace(tokenizer=_minicpmo45_tokenizer_stub())
    load_calls = []

    class _AutoImageProcessor:
        @staticmethod
        def register(config_class, image_processor_class):
            del image_processor_class
            if isinstance(config_class, str):
                raise AttributeError("'str' object has no attribute '__module__'")

    class _AutoProcessor:
        @classmethod
        def from_pretrained(cls, model_path, *, trust_remote_code):
            load_calls.append((model_path, trust_remote_code))
            _AutoImageProcessor.register("MiniCPMVImageProcessor", object)
            return processor

    monkeypatch.setitem(
        sys.modules,
        "transformers",
        SimpleNamespace(
            AutoImageProcessor=_AutoImageProcessor,
            AutoProcessor=_AutoProcessor,
        ),
    )

    loaded = MiniCPMO45Stage0DuplexRuntime._load_processor_from_path(model_id)

    assert loaded is processor
    assert load_calls == [(model_id, True)]
    with pytest.raises(AttributeError, match="__module__"):
        _AutoImageProcessor.register("MiniCPMVImageProcessor", object)


def test_minicpmo_stage0_processor_load_preserves_original_exception(monkeypatch):
    import sys

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
    )

    original = ValueError("invalid processor config")

    class _AutoProcessor:
        @classmethod
        def from_pretrained(cls, model_path, *, trust_remote_code):
            del cls, model_path, trust_remote_code
            raise original

    monkeypatch.setitem(
        sys.modules,
        "transformers",
        _transformers_processor_stub(_AutoProcessor),
    )

    with pytest.raises(RuntimeError, match="Failed to load MiniCPM-o duplex processor") as exc_info:
        MiniCPMO45Stage0DuplexRuntime._load_processor_from_path("openbmb/MiniCPM-o-4_5")

    assert exc_info.value.__cause__ is original


def test_minicpmo_stage0_processor_requires_tokenizer(monkeypatch):
    import sys

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
    )

    class _AutoProcessor:
        @classmethod
        def from_pretrained(cls, model_path, *, trust_remote_code):
            del cls, model_path, trust_remote_code
            return object()

    monkeypatch.setitem(
        sys.modules,
        "transformers",
        _transformers_processor_stub(_AutoProcessor),
    )

    with pytest.raises(RuntimeError, match="does not expose a tokenizer"):
        MiniCPMO45Stage0DuplexRuntime._load_processor_from_path("openbmb/MiniCPM-o-4_5")


def test_minicpmo_stage0_loaded_processor_validates_special_tokens(monkeypatch):
    import sys

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
    )

    processor = SimpleNamespace(tokenizer=_minicpmo45_tokenizer_stub())

    class _AutoProcessor:
        @classmethod
        def from_pretrained(cls, model_path, *, trust_remote_code):
            del cls, model_path, trust_remote_code
            return processor

    monkeypatch.setitem(
        sys.modules,
        "transformers",
        _transformers_processor_stub(_AutoProcessor),
    )

    runtime = MiniCPMO45Stage0DuplexRuntime(
        SimpleNamespace(),
        model_path="openbmb/MiniCPM-o-4_5",
        device="cpu",
    )

    assert runtime.processor is processor
    assert set(runtime._special_token_ids()) == {
        "unit_token_id",
        "unit_end_token_id",
        "listen_token_id",
        "speak_token_id",
        "tts_bos_token_id",
        "tts_eos_token_id",
        "tts_pad_token_id",
        "chunk_eos_token_id",
        "chunk_tts_eos_token_id",
        "turn_eos_token_id",
    }


def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    class _StageModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.embed = torch.nn.Embedding(256, 2)

        def get_input_embeddings(self):
            return self.embed

        def get_audio_hidden_states(self, _data):
            return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = _StageModel()
    runtime.thinker = runtime.stage_model
    runtime.tokenizer = SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: {
            "<unit>": 1,
            "</unit>": 2,
            "<|listen|>": 3,
            "<|speak|>": 4,
            "<|tts_bos|>": 5,
            "<|tts_eos|>": 6,
            "<|tts_pad|>": 7,
            "<|chunk_eos|>": 8,
            "<|chunk_tts_eos|>": 9,
            "<|turn_eos|>": 10,
            "<|audio|>": 11,
        }.get(token, 0),
        encode=lambda text, add_special_tokens=False: [201, 5],
    )
    runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4)
    runtime.device = "cpu"
    runtime._init_token_ids()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-data-plane-prefill")

    # Official duplex format: each unit is <unit> + audio embeddings with no
    # per-chunk assistant header or <|tts_bos|> boundary. Decoding starts right
    # after the audio so the first sampled token is the listen/speak decision.
    result = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)

    assert result["success"] is True
    assert result["input_token_ids"] == [1, 11]
    assert result["prompt_suffix_len"] == 0

    # Subsequent units must close the previous unit with </unit> first,
    # mirroring the official finalize_unit() feed.
    result = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=2)

    assert result["success"] is True
    assert result["input_token_ids"] == [2, 1, 11]
    assert result["prompt_suffix_len"] == 0


@pytest.mark.parametrize("mode", [None, "off", "basic", "context"])
def test_minicpmo_stage0_history_collection_starts_with_session_config(mode):
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState

    runtime = _stage0_vision_runtime()
    state = _MiniCPMO45Stage0SessionState(session_id="history-lifecycle")
    config = {} if mode is None else {"duplex_window_config": {"sliding_window_mode": mode}}
    runtime._prepare_session_context(state, {}, runtime_config=config)
    for seq in range(1, 101):
        runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=seq)
    enabled = mode in {"basic", "context"}
    assert state.window_enabled is enabled
    assert len(state.window_units) == (99 if enabled else 0)
    assert (state.pending_window_unit is not None) is enabled
    assert state.pending_window_generated_tokens == []


@pytest.mark.parametrize("mode", ["basic", "context"])
@pytest.mark.parametrize("buffered", [False, True])
def test_minicpmo_stage0_multi_chunk_append_matches_scheduler_window(mode, buffered):
    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState

    runtime = _stage0_vision_runtime()
    state = _MiniCPMO45Stage0SessionState(session_id="multi-chunk-window", window_enabled=True)
    if buffered:
        # Keep the tiny fake processor's initial chunk unpadded to model carry.
        runtime._pad_first_audio_chunk_if_needed = lambda *args: None
        assert runtime._stage_prefill_embeddings_only(state, np.zeros(2, dtype=np.float32), seq=0)["success"] is False
    first = runtime._stage_prefill_embeddings_only(
        state, np.zeros(6 if buffered else 8, dtype=np.float32), seq=1, final=buffered
    )
    assert first["input_token_ids"] == [1, 11, 2, 1, 11]
    assert state.window_units == []
    assert state.pending_window_unit.token_ids == first["input_token_ids"]
    session = SimpleNamespace(num_prompt_tokens=5)
    replaced = False
    for seq in (2, 3):
        info = {
            "duplex": {
                "data_plane": True,
                "seq": seq,
                "runtime_config": {
                    "duplex_window_config": {
                        "sliding_window_mode": mode,
                        "basic_window_high_tokens": 1,
                        "basic_window_low_tokens": 2,
                        "context_max_units": 1,
                        "context_previous_max_tokens": 16,
                    },
                    "duplex_window_previous_marker_token_ids": [201, 5],
                },
            }
        }
        update = SimpleNamespace(prompt_token_ids=[3, 2, 1, 11], model_intermediate_buffer=info)
        planned = OmniARScheduler._prepare_minicpmo45_stage0_window(
            session, update, segment_output_ids=[44], completed_terminator=3
        )
        result = runtime._stage_prefill_embeddings_only(
            state, np.zeros(4, dtype=np.float32), seq=seq, stage0_window=info["duplex"]["stage0_window"]
        )
        assert result["stage0_window_replaced"] is planned
        if planned:
            assert result["num_input_tokens"] == len(update.prompt_token_ids)
            replaced = True
            break
        session.num_prompt_tokens += 1 + len(update.prompt_token_ids)
    assert replaced


def test_minicpmo_stage0_basic_window_rebuilds_from_retained_units():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )

    runtime = _stage0_vision_runtime()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-basic-window", window_enabled=True)
    first = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)
    assert first["input_token_ids"] == [1, 11]
    state.pending_window_generated_tokens.append(5)
    state.pending_terminator_token = 3

    rebuilt = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        seq=2,
        stage0_window={
            "replace": True,
            "mode": "basic",
            "drop_units": 1,
            "replacement_prompt_len": 2,
        },
    )

    assert rebuilt["stage0_window_replaced"] is True
    assert rebuilt["input_token_ids"] == [1, 11]
    assert state.window_units == []


def test_minicpmo_stage0_context_window_inserts_previous_before_suffix():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )

    runtime = _stage0_vision_runtime()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-context-window", window_enabled=True)
    prefix = runtime._embed_token(200)
    suffix = runtime._embed_token(202)
    state.context_embeds = [prefix, suffix]
    state.context_token_ids = [200, 202]
    state.context_prefix_embeds = [prefix]
    state.context_prefix_token_ids = [200]
    state.context_suffix_embeds = [suffix]
    state.context_suffix_token_ids = [202]
    first = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)
    assert first["input_token_ids"] == [200, 202, 1, 11]
    state.pending_window_generated_tokens.append(42)
    state.pending_terminator_token = 3

    rebuilt = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        seq=2,
        stage0_window={
            "replace": True,
            "mode": "context",
            "drop_units": 1,
            "previous_token_ids": [42],
            "previous_marker_token_ids": [201, 5],
            "replacement_prompt_len": 7,
        },
    )

    assert rebuilt["input_token_ids"] == [200, 201, 5, 42, 202, 1, 11]
    assert rebuilt["num_input_tokens"] == 7


@pytest.mark.parametrize("pending_terminator", [None, 3, 99])
def test_minicpmo_stage0_window_uses_accepted_output_not_async_sampler_history(pending_terminator):
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )

    runtime = _stage0_vision_runtime()
    runtime._stage_audio_embeddings = lambda *args, **kwargs: torch.zeros((10, 2))
    state = _MiniCPMO45Stage0SessionState(session_id="sid-accepted-window", window_enabled=True)
    state.context_prefix_token_ids = [200] * 7
    state.context_suffix_token_ids = [202]
    state.context_token_ids = [*state.context_prefix_token_ids, 202]
    state.context_prefix_embeds = [runtime._embed_token(token_id) for token_id in state.context_prefix_token_ids]
    state.context_suffix_embeds = [runtime._embed_token(202)]
    state.context_embeds = [*state.context_prefix_embeds, *state.context_suffix_embeds]
    runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)
    for seq in (2, 3):
        state.pending_window_generated_tokens = [3, 40, 41]
        state.pending_terminator_token = pending_terminator
        plan = {"completed_token_ids": [], "completed_terminator_token_id": 3}
        if seq == 3:
            plan.update(replace=True, mode="context", drop_units=1, replacement_prompt_len=32)
        result = runtime._stage_prefill_embeddings_only(
            state, np.zeros(4, dtype=np.float32), seq=seq, stage0_window=plan
        )

    assert result["num_input_tokens"] == 32
    assert [len(unit.token_ids) for unit in state.window_units] == [13]
    assert len(state.pending_window_unit.token_ids) == 11
    assert 40 not in result["input_token_ids"]
    assert 41 not in result["input_token_ids"]


def _stage0_vision_runtime():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
    )

    class _StageModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.embed = torch.nn.Embedding(256, 2)

        def get_input_embeddings(self):
            return self.embed

        def get_audio_hidden_states(self, _data):
            return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = _StageModel()
    runtime.thinker = runtime.stage_model
    runtime.tokenizer = SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: {
            "<unit>": 1,
            "</unit>": 2,
            "<|listen|>": 3,
            "<|speak|>": 4,
            "<|tts_bos|>": 5,
            "<|tts_eos|>": 6,
            "<|tts_pad|>": 7,
            "<|chunk_eos|>": 8,
            "<|chunk_tts_eos|>": 9,
            "<|turn_eos|>": 10,
            "<|audio|>": 11,
            "<image>": 12,
            "</image>": 13,
            "<slice>": 14,
            "</slice>": 15,
        }.get(token, 0),
        encode=lambda text, add_special_tokens=False: [201, 5],
    )
    runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4)
    runtime.device = "cpu"
    runtime._init_token_ids()
    # One 2-row block per frame keeps the expected token_ids readable; the real
    # resampler emits VISION_EMBEDS_PER_FRAME rows per slice.
    runtime._stage_vision_embeddings = lambda frames: [
        [torch.full((2, 2), float(index), dtype=torch.float32)] for index in range(len(frames))
    ]
    return runtime


def _minicpmo_seeded_silence_sampling_case(initial_user_text: str | None, *, force_listen_count: int = 0):
    from vllm_omni.engine.duplex.contracts import DuplexFence
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import build_duplex_data_plane_prompt
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState

    runtime = _stage0_vision_runtime()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-seeded-text")
    runtime_config = {"initial_user_text": initial_user_text}
    runtime._prepare_session_context(state, {"instructions": "Speak exactly."}, runtime_config=runtime_config)
    prefill = runtime._stage_prefill_embeddings_only(
        state, np.zeros(4, dtype=np.float32), epoch=0, seq=1, is_speech=False
    )
    assert prefill["success"] is True
    prompt = build_duplex_data_plane_prompt(
        request_id="req-seeded-text",
        fence=DuplexFence(state.session_id),
        session_config={"extra_body": {"force_listen_count": force_listen_count}},
        runtime_config=runtime_config,
        seq=1,
        turn_seq=1,
        payload={"type": "audio", "is_speech": False},
        final=False,
    )
    payload = prompt["model_intermediate_buffer"]["duplex"]["payload"]
    model, row = _minicpmo_duplex_policy_case(state, payload)
    model._minicpmo45_native_duplex_token_ids_cache = runtime._special_token_ids()
    return runtime, state, model, row


@pytest.mark.parametrize("initial_user_text", [None, "", "Get the trust fund to the bank early."])
@pytest.mark.parametrize("force_listen_count", [0, 1])
def test_minicpmo_seeded_text_reaches_sampling_on_silence(initial_user_text, force_listen_count):
    """A seeded user turn must reach the model despite its silent audio clock."""
    runtime, state, model, row = _minicpmo_seeded_silence_sampling_case(
        initial_user_text, force_listen_count=force_listen_count
    )
    logits = torch.zeros((1, 32), dtype=torch.float32)
    logits[0, 20] = 10.0
    original_logits = logits.clone()

    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))

    if initial_user_text and not force_listen_count:
        assert torch.equal(logits, original_logits)
        assert state.pending_speech_context is True
    else:
        assert logits[0, runtime.listen_token_id] == 0
        assert torch.isfinite(logits).sum() == 1


def test_minicpmo_seeded_text_is_consumed_before_post_turn_silence():
    """The opening text allows one answer; silence after its end stays quiet."""
    from dataclasses import replace

    runtime, state, model, row = _minicpmo_seeded_silence_sampling_case("Read this sentence.")
    token_ids = runtime._special_token_ids()
    logits = torch.zeros((1, 32), dtype=torch.float32)
    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
    assert torch.isfinite(logits).all()

    model._record_minicpmo45_duplex_terminator(0, runtime.tts_bos_token_id, token_ids)
    assert state.pending_speech_context is True
    model._record_minicpmo45_duplex_terminator(0, 20, token_ids)
    assert state.pending_speech_context is False
    model._record_minicpmo45_duplex_terminator(0, runtime.turn_eos_token_id, token_ids)
    assert state.current_turn_ended is True

    prefill = runtime._stage_prefill_embeddings_only(
        state, np.zeros(4, dtype=np.float32), epoch=0, seq=2, is_speech=False
    )
    assert prefill["success"] is True
    next_logits = torch.zeros((1, 32), dtype=torch.float32)
    model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (replace(row, seq=row.seq + 1),))
    assert next_logits[0, runtime.listen_token_id] == 0
    assert torch.isfinite(next_logits).sum() == 1


def test_minicpmo_seeded_text_survives_explicit_initial_listen():
    """The force-listen prefix delays the seed without losing its answer."""
    from dataclasses import replace

    runtime, state, model, row = _minicpmo_seeded_silence_sampling_case("Read this sentence.", force_listen_count=1)
    token_ids = runtime._special_token_ids()
    logits = torch.zeros((1, 32), dtype=torch.float32)
    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
    assert torch.isfinite(logits).sum() == 1
    model._record_minicpmo45_duplex_terminator(0, runtime.listen_token_id, token_ids)

    prefill = runtime._stage_prefill_embeddings_only(
        state, np.zeros(4, dtype=np.float32), epoch=0, seq=2, is_speech=False
    )
    assert prefill["success"] is True
    next_logits = torch.zeros((1, 32), dtype=torch.float32)
    next_row = replace(row, seq=row.seq + 1, payload={"is_speech": False, "force_listen": False})
    model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (next_row,))
    assert torch.isfinite(next_logits).all()
    assert state.pending_speech_context is True


@pytest.mark.parametrize("initial_text", [None, "", "Please say hello.", "请说你好。"])
@pytest.mark.parametrize("speech_envelope", [False, True])
def test_minicpmo_seeded_context_keeps_assistant_turn_open_until_eos(initial_text, speech_envelope, mocker):
    from vllm.v1.sample.metadata import SamplingMetadata

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState

    runtime = _stage0_vision_runtime()
    state = _MiniCPMO45Stage0SessionState(session_id="seeded-context")
    runtime._prepare_session_context(state, {}, runtime_config={"initial_user_text": initial_text})
    result = runtime._stage_prefill_embeddings_only(
        state, np.zeros(4, dtype=np.float32), epoch=0, seq=1, is_speech=False
    )
    assert result["success"] is True
    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False})
    logits = torch.zeros((1, 16), dtype=torch.float32)
    logits[0, 10] = 20.0
    original_logits = logits.clone()
    sampling_metadata = mocker.Mock(spec=SamplingMetadata)
    model.prepare_duplex_sampling(logits, sampling_metadata, (row,))
    token_ids = model._minicpmo45_native_duplex_token_ids_cache
    sampled = model._finalize_minicpmo45_native_duplex_sample(0, 7, token_ids)

    if not initial_text:
        assert sampled == 7
        assert logits[0, 7].item() == 0.0
        assert torch.isneginf(logits[0, :7]).all()
        assert torch.isneginf(logits[0, 8:]).all()
        return

    assert torch.equal(logits, original_logits)
    assert sampled == 8
    if speech_envelope:
        model._record_minicpmo45_duplex_terminator(0, 8, token_ids)
        assert state.current_turn_ended is False
        assert state.pending_speech_context is True
    model._record_minicpmo45_duplex_terminator(0, 10, token_ids)
    assert state.pending_speech_context is False
    model._record_minicpmo45_duplex_terminator(0, 9, token_ids)

    # A later silent unit must not re-admit the seed after the turn ended.
    result = runtime._stage_prefill_embeddings_only(
        state, np.zeros(4, dtype=np.float32), epoch=0, seq=2, is_speech=False
    )
    assert result["success"] is True
    logits = original_logits.clone()
    model.prepare_duplex_sampling(logits, sampling_metadata, (replace(row, seq=4),))
    assert logits[0, 7].item() == 0.0
    assert torch.isneginf(logits[0, :7]).all()
    assert torch.isneginf(logits[0, 8:]).all()
    assert model._finalize_minicpmo45_native_duplex_sample(0, 7, token_ids) == 7


def test_minicpmo_stage0_puts_every_frame_of_an_append_in_one_unit():
    """Official streaming_prefill feeds the whole frame_list into one unit.

    Stacked pairs used to be spread one frame per unit, so a two-frame append
    that closed a single unit left a frame over, failed the prefill, and still
    advanced the session state.
    """
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )

    runtime = _stage0_vision_runtime()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-stacked-frames")

    result = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        video_frames=["frame-prev", "frame-cur"],
        seq=1,
    )

    assert result["success"] is True
    # <unit>, then each frame as <image> + embeds + </image>, then the audio.
    assert result["input_token_ids"] == [1, 12, 0, 0, 13, 12, 0, 0, 13, 11]
    assert state.audio_chunk_idx == 1


def test_minicpmo_stage0_wraps_hd_slices_after_the_source_image():
    """Official stacked pair is max_slice_nums=[2, 1]: source + slices, then audio."""
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )

    runtime = _stage0_vision_runtime()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-hd-slices")
    runtime._stage_vision_embeddings = lambda frames: [
        [torch.ones((2, 2)), torch.full((2, 2), 2.0), torch.full((2, 2), 3.0)],
        [torch.full((2, 2), 4.0)],
    ]

    result = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        video_frames=["base", "stack"],
        seq=1,
    )

    assert result["success"] is True
    # <unit> <image> src </image> <slice> p0 </slice> <slice> p1 </slice>
    # <image> stack </image> audio
    assert result["input_token_ids"] == [
        1,
        12,
        0,
        0,
        13,
        14,
        0,
        0,
        15,
        14,
        0,
        0,
        15,
        12,
        0,
        0,
        13,
        11,
    ]


def test_minicpmo_stage0_reports_frames_that_close_no_unit_instead_of_dropping_them():
    """A frame only binds to the unit its append closes.

    Frames are not carried across appends, so one arriving early used to vanish
    under a generic "audio not enough" buffering result.
    """
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )

    runtime = _stage0_vision_runtime()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-early-frame")

    assert runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)["success"] is True

    early = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(2, dtype=np.float32),
        video_frames=["frame-cur"],
        seq=2,
    )

    assert early["success"] is False
    assert "closes no model unit" in early["reason"]
    # The half-filled append must not advance the unit counter.
    assert state.audio_chunk_idx == 1


def test_minicpmo_stage0_speech_append_sets_pending_context_once():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    class _StageModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.embed = torch.nn.Embedding(256, 2)

        def get_input_embeddings(self):
            return self.embed

        def get_audio_hidden_states(self, _data):
            return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = _StageModel()
    runtime.thinker = runtime.stage_model
    runtime.tokenizer = SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: {
            "<unit>": 1,
            "</unit>": 2,
            "<|listen|>": 3,
            "<|speak|>": 4,
            "<|tts_bos|>": 5,
            "<|tts_eos|>": 6,
            "<|tts_pad|>": 7,
            "<|chunk_eos|>": 8,
            "<|chunk_tts_eos|>": 9,
            "<|turn_eos|>": 10,
            "<|audio|>": 11,
        }.get(token, 0),
        encode=lambda text, add_special_tokens=False: [],
    )
    runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4)
    runtime.device = "cpu"
    runtime._init_token_ids()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-speech-append")

    result = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        epoch=0,
        seq=1,
        is_speech=True,
    )

    assert result["success"] is True
    assert state.pending_speech_context is True
    assert state.pending_speech_append_identity == (0, 1)

    state.pending_speech_context = False
    cached = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        epoch=0,
        seq=1,
        is_speech=True,
    )

    assert cached["success"] is True
    assert state.pending_speech_context is False
    assert state.pending_speech_append_identity == (0, 1)

    next_seq = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        epoch=0,
        seq=2,
        is_speech=True,
    )

    assert next_seq["success"] is True
    assert state.pending_speech_context is True
    assert state.pending_speech_append_identity == (0, 2)

    state.pending_speech_context = False
    next_epoch = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        epoch=1,
        seq=1,
        is_speech=True,
    )

    assert next_epoch["success"] is True
    assert state.pending_speech_context is True
    assert state.pending_speech_append_identity == (1, 1)

    state.pending_speech_context = False
    silence = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        epoch=1,
        seq=2,
        is_speech=False,
    )

    assert silence["success"] is True
    assert state.pending_speech_context is False
    assert state.pending_speech_append_identity == (1, 1)


def test_minicpmo_stage0_failed_append_does_not_set_pending_speech_context():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    class _StageModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.embed = torch.nn.Embedding(256, 2)

        def get_input_embeddings(self):
            return self.embed

        def get_audio_hidden_states(self, _data):
            return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = _StageModel()
    runtime.thinker = runtime.stage_model
    runtime.tokenizer = SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: {
            "<unit>": 1,
            "</unit>": 2,
            "<|listen|>": 3,
            "<|speak|>": 4,
            "<|tts_bos|>": 5,
            "<|tts_eos|>": 6,
            "<|tts_pad|>": 7,
            "<|chunk_eos|>": 8,
            "<|chunk_tts_eos|>": 9,
            "<|turn_eos|>": 10,
            "<|audio|>": 11,
        }.get(token, 0),
        encode=lambda text, add_special_tokens=False: [],
    )
    runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4)
    runtime.device = "cpu"
    runtime._init_token_ids()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-failed-speech")

    result = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(0, dtype=np.float32),
        epoch=0,
        seq=1,
        is_speech=True,
    )

    assert result["success"] is False
    assert state.pending_speech_context is False
    assert state.pending_speech_append_identity is None


def test_minicpmo_stage0_streaming_processor_is_isolated_per_session():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    class _Mel:
        def __init__(self):
            self.counter = 0

    class _Processor:
        def __init__(self):
            self._streaming_mel_processor = _Mel()

        def set_streaming_mode(self, **_kwargs):
            return None

        def reset_streaming(self):
            self._streaming_mel_processor.counter = 0

        def process_audio_streaming(self, _audio, **_kwargs):
            self._streaming_mel_processor.counter += 1
            return self._streaming_mel_processor.counter

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.processor = _Processor()
    runtime.stage_model = SimpleNamespace()
    runtime.thinker = runtime.stage_model
    a = _MiniCPMO45Stage0SessionState(session_id="a")
    b = _MiniCPMO45Stage0SessionState(session_id="b")
    processor_a = runtime._configure_streaming_processor(a)
    processor_b = runtime._configure_streaming_processor(b)

    observed = [
        runtime._process_streaming_audio([], 0, processor=processor_a),
        runtime._process_streaming_audio([], 0, processor=processor_b),
        runtime._process_streaming_audio([], 1, processor=processor_a),
        runtime._process_streaming_audio([], 1, processor=processor_b),
    ]

    assert observed == [1, 1, 2, 2]
    assert processor_a is not processor_b
    assert runtime.processor._streaming_mel_processor.counter == 0


def test_minicpmo_stage0_data_plane_next_append_reinjects_previous_listen():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    class _StageModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.embed = torch.nn.Embedding(256, 2)

        def get_input_embeddings(self):
            return self.embed

        def get_audio_hidden_states(self, data):
            return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = _StageModel()
    runtime.thinker = runtime.stage_model
    runtime.tokenizer = SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: {
            "<unit>": 1,
            "</unit>": 2,
            "<|listen|>": 3,
            "<|speak|>": 4,
            "<|tts_bos|>": 5,
            "<|tts_eos|>": 6,
            "<|tts_pad|>": 7,
            "<|chunk_eos|>": 8,
            "<|chunk_tts_eos|>": 9,
            "<|turn_eos|>": 10,
            "<|audio|>": 11,
        }.get(token, 0),
        encode=lambda text, add_special_tokens=False: [],
    )
    runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4)
    runtime.device = "cpu"
    runtime._init_token_ids()
    state = _MiniCPMO45Stage0SessionState(
        session_id="sid-new-speech-prefill",
        audio_chunk_idx=1,
        pending_terminator_token=3,
        last_terminator_token=3,
        current_turn_ended=True,
    )

    result = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        seq=2,
    )

    assert result["success"] is True
    assert result["input_token_ids"] == [3, 2, 1, 11]
    assert state.pending_terminator_token is None
    assert state.last_terminator_token == 3
    assert state.current_turn_ended is True


def test_minicpmo_stage0_data_plane_turn_eos_closes_previous_unit():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    class _StageModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.embed = torch.nn.Embedding(256, 2)

        def get_input_embeddings(self):
            return self.embed

        def get_audio_hidden_states(self, data):
            return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = _StageModel()
    runtime.thinker = runtime.stage_model
    runtime.tokenizer = SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: {
            "<unit>": 1,
            "</unit>": 2,
            "<|listen|>": 3,
            "<|speak|>": 4,
            "<|tts_bos|>": 5,
            "<|tts_eos|>": 6,
            "<|tts_pad|>": 7,
            "<|chunk_eos|>": 8,
            "<|chunk_tts_eos|>": 9,
            "<|turn_eos|>": 10,
            "<|audio|>": 11,
        }.get(token, 0),
        encode=lambda text, add_special_tokens=False: [],
    )
    runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4)
    runtime.device = "cpu"
    runtime._init_token_ids()
    state = _MiniCPMO45Stage0SessionState(
        session_id="sid-new-user-turn-prefill",
        audio_chunk_idx=1,
        pending_terminator_token=10,
        last_terminator_token=10,
        current_turn_ended=True,
    )

    result = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        seq=2,
    )

    assert result["success"] is True
    assert result["input_token_ids"] == [10, 2, 1, 11]
    assert state.pending_terminator_token is None
    assert state.last_terminator_token == 10
    assert state.current_turn_ended is True


def test_minicpmo_stage0_data_plane_model_owned_turn_boundary_preserves_audio_cache():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    stale_cache = object()

    class _StageModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.embed = torch.nn.Embedding(256, 2)
            self.audio_past_key_values = stale_cache

        def get_input_embeddings(self):
            return self.embed

        def get_audio_hidden_states(self, data):
            return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = _StageModel()
    runtime.thinker = runtime.stage_model
    runtime.tokenizer = SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: {
            "<unit>": 1,
            "</unit>": 2,
            "<|listen|>": 3,
            "<|speak|>": 4,
            "<|tts_bos|>": 5,
            "<|tts_eos|>": 6,
            "<|tts_pad|>": 7,
            "<|chunk_eos|>": 8,
            "<|chunk_tts_eos|>": 9,
            "<|turn_eos|>": 10,
            "<|audio|>": 11,
        }.get(token, 0),
        encode=lambda text, add_special_tokens=False: [],
    )
    runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4)
    runtime.device = "cpu"
    runtime._init_token_ids()
    state = _MiniCPMO45Stage0SessionState(
        session_id="sid-new-user-turn-audio-cache",
        audio_chunk_idx=1,
        audio_past_key_values=stale_cache,
        pending_terminator_token=8,
        last_terminator_token=8,
        current_turn_ended=True,
    )

    result = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(4, dtype=np.float32),
        seq=2,
    )

    assert result["success"] is True
    assert state.audio_past_key_values is stale_cache
    assert runtime.thinker.audio_past_key_values is stale_cache


def test_minicpmo_stage0_data_plane_final_first_chunk_does_not_add_silence_unit():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    class _StageModel(torch.nn.Module):
        first_chunk_ms = 10
        sample_rate = 1000

        def __init__(self):
            super().__init__()
            self.seen_audio = None
            self.embed = torch.nn.Embedding(256, 2)

        def get_input_embeddings(self):
            return self.embed

        def get_audio_hidden_states(self, data):
            self.seen_audio = np.asarray(data["audio_features"], dtype=np.float32)
            return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]

    class _MelProcessor:
        sample_rate = 1000

        def get_config(self):
            return {"effective_first_chunk_ms": 10}

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = _StageModel()
    runtime.thinker = runtime.stage_model
    runtime.tokenizer = SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: {
            "<unit>": 1,
            "</unit>": 2,
            "<|listen|>": 3,
            "<|speak|>": 4,
            "<|tts_bos|>": 5,
            "<|tts_eos|>": 6,
            "<|tts_pad|>": 7,
            "<|chunk_eos|>": 8,
            "<|chunk_tts_eos|>": 9,
            "<|turn_eos|>": 10,
            "<|audio|>": 11,
        }.get(token, 0),
        encode=lambda text, add_special_tokens=False: [],
    )
    runtime.processor = SimpleNamespace(
        _streaming_mel_processor=_MelProcessor(),
        get_streaming_chunk_size=lambda: 10,
    )
    runtime.device = "cpu"
    runtime._init_token_ids()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-first-chunk-padding")

    result = runtime._stage_prefill_embeddings_only(
        state,
        np.arange(8, dtype=np.float32),
        seq=1,
        final=True,
    )

    assert result["success"] is True
    assert result["input_token_ids"] == [1, 11]
    assert runtime.stage_model.seen_audio is not None
    np.testing.assert_allclose(
        runtime.stage_model.seen_audio.reshape(-1),
        np.array([0, 0, 0, 1, 2, 3, 4, 5, 6, 7], dtype=np.float32),
    )


def test_minicpmo_stage0_final_does_not_promote_first_chunk_alignment_tail_to_unit():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    class _StageModel(torch.nn.Module):
        first_chunk_ms = 10
        chunk_ms = 8
        sample_rate = 1000

        def __init__(self):
            super().__init__()
            self.embed = torch.nn.Embedding(256, 2)

        def get_input_embeddings(self):
            return self.embed

        def get_audio_hidden_states(self, _data):
            return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]

    class _MelProcessor:
        sample_rate = 1000
        is_first = True

        def get_config(self):
            return {"effective_first_chunk_ms": 9}

    class _Processor:
        def __init__(self):
            self._streaming_mel_processor = _MelProcessor()

        def get_streaming_chunk_size(self):
            return 9 if self._streaming_mel_processor.is_first else 8

        def process_audio_streaming(self, audio, **_kwargs):
            self._streaming_mel_processor.is_first = False
            return {"audio_features": audio, "audio_feature_lens": [[len(audio)]]}

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = _StageModel()
    runtime.thinker = runtime.stage_model
    runtime.tokenizer = SimpleNamespace(
        unk_token_id=0,
        convert_tokens_to_ids=lambda token: {
            "<unit>": 1,
            "</unit>": 2,
            "<|listen|>": 3,
            "<|speak|>": 4,
            "<|tts_bos|>": 5,
            "<|tts_eos|>": 6,
            "<|tts_pad|>": 7,
            "<|chunk_eos|>": 8,
            "<|chunk_tts_eos|>": 9,
            "<|turn_eos|>": 10,
            "<|audio|>": 11,
        }.get(token, 0),
        encode=lambda text, add_special_tokens=False: [],
    )
    runtime.processor = _Processor()
    runtime.device = "cpu"
    runtime._init_token_ids()
    state = _MiniCPMO45Stage0SessionState(session_id="sid-first-chunk-alignment-tail")

    first = runtime._stage_prefill_embeddings_only(state, np.zeros(8, dtype=np.float32), seq=1)
    second = runtime._stage_prefill_embeddings_only(state, np.zeros(8, dtype=np.float32), seq=2)
    final = runtime._stage_prefill_embeddings_only(
        state,
        np.zeros(8, dtype=np.float32),
        seq=3,
        final=True,
    )

    assert first["input_token_ids"].count(1) == 1
    assert second["input_token_ids"].count(1) == 1
    assert final["input_token_ids"].count(1) == 1
    assert state.audio_chunk_idx == 3
    assert len(state.audio_buffer) == 1


def test_minicpmo_stage0_runtime_uses_loaded_vllm_embed_tokens_when_get_input_embeddings_is_broken():
    import torch

    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
    )

    class _Embed(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.weight = torch.nn.Parameter(torch.ones(128, 2))
            self.calls = []

        def forward(self, input_ids):
            ids = input_ids.reshape(-1).tolist()
            self.calls.append(ids)
            return torch.tensor([[float(i), 0.0] for i in ids], dtype=torch.float32)

    class _Thinker:
        def __init__(self):
            self.llm = SimpleNamespace(model=SimpleNamespace(embed_tokens=_Embed()))

        def get_input_embeddings(self, input_ids, multimodal_embeddings=None):
            raise AttributeError("'Qwen3ForCausalLM' object has no attribute 'get_input_embeddings'")

    thinker = _Thinker()
    stage_model = SimpleNamespace(model_stage="llm", thinker=thinker, processor=None)
    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = stage_model
    runtime.thinker = thinker
    runtime.device = "cpu"

    embeds = runtime._embed_token(11)

    assert embeds.shape == (1, 2)
    assert thinker.llm.model.embed_tokens.calls == [[11]]


def test_minicpmo_remote_config_patch_handles_nested_and_dict_configs():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.compat import (
        patch_minicpmo_remote_config,
    )

    nested = SimpleNamespace(base_model_tp_plan=None)
    config = SimpleNamespace(
        base_model_tp_plan=None,
        text_config=nested,
        tts_config={},
    )

    patch_minicpmo_remote_config(config)

    assert config.base_model_tp_plan == {}
    assert nested.base_model_tp_plan == {}
    assert config.tts_config["top_p"] == 0.8
    assert config.tts_config["top_k"] == 100
    assert config.tts_config["temperature"] == 0.8
    assert config.tts_config["repetition_penalty"] == 1.05


def test_minicpmo_stage0_short_audio_buffers_without_context_mutation():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    class _Processor:
        def get_streaming_chunk_size(self):
            return 16000

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.stage_model = SimpleNamespace()
    runtime.thinker = SimpleNamespace()
    runtime.processor = _Processor()
    runtime._require_special_token_ids = lambda: None
    state = _MiniCPMO45Stage0SessionState(session_id="sid")

    result = runtime._stage_prefill_embeddings_only(state, np.zeros(1600, dtype=np.float32))

    assert result["success"] is False
    assert result["reason"]
    assert len(state.audio_buffer) >= 1600
    assert state.context_embeds == []


def test_minicpmo_stage0_native_sampler_penalizes_repeated_text_token():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    vocab_size = 151723
    repeated = 198
    alternative = 1234
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, repeated] = 20.0
    logits[0, alternative] = 19.5
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[repeated] * 8],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[alternative]]


def test_minicpmo_stage0_native_sampler_penalizes_text_from_prior_chunk():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    session_key = "sid-cross-chunk-repetition"
    state = _MiniCPMO45Stage0SessionState(session_id=session_key)
    state.generated_tokens = [198] * 8
    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state})
    model._minicpmo45_duplex_row_sessions = {0: session_key}

    vocab_size = 151723
    repeated = 198
    alternative = 1234
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, repeated] = 20.0
    logits[0, alternative] = 19.5
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[alternative]]


def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    session_key = "sid-negative-logit-repetition"
    state = _MiniCPMO45Stage0SessionState(session_id=session_key)
    repeated = 198
    alternative = 1234
    state.generated_tokens = [repeated]
    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state})
    model._minicpmo45_duplex_row_sessions = {0: session_key}

    logits = torch.full((1, 151723), -100.0)
    logits[0, repeated] = -1.0
    logits[0, alternative] = -0.97
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[repeated]]


def test_minicpmo_stage0_records_bounded_model_policy_history():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    session_key = "sid-text-history"
    state = _MiniCPMO45Stage0SessionState(session_id=session_key)
    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state})
    model._minicpmo45_duplex_row_sessions = {0: session_key}
    token_ids = {
        "unit_token_id": 1,
        "unit_end_token_id": 2,
        "listen_token_id": 3,
        "speak_token_id": 4,
        "tts_bos_token_id": 5,
        "tts_eos_token_id": 6,
        "tts_pad_token_id": 7,
        "chunk_eos_token_id": 8,
        "chunk_tts_eos_token_id": 9,
        "turn_eos_token_id": 10,
    }

    for sampled in range(1000, 1520):
        model._record_minicpmo45_duplex_generation_token(0, sampled)

    expected_history = list(range(1008, 1520))
    assert state.generated_tokens == expected_history

    for sampled in token_ids.values():
        model._record_minicpmo45_duplex_generation_token(0, sampled)

    assert state.generated_tokens == [*expected_history[len(token_ids) :], *token_ids.values()]


def test_minicpmo_stage0_native_sampler_does_not_override_model_at_punctuation():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        text = {
            200: "我",
            201: "喜",
            202: "欢",
            203: "。",
        }

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

        def decode(self, ids, skip_special_tokens=True):
            del skip_special_tokens
            return "".join(self.text.get(int(token_id), "") for token_id in ids)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    model.min_new_speak_tokens_before_chunk_boundary = 4
    model.max_new_speak_tokens_per_chunk = 64
    vocab_size = 151723
    alternative = 1234
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, alternative] = 20.0
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[151706, 200, 201, 202, 203]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[alternative]]


def test_minicpmo_stage0_native_sampler_does_not_cut_before_natural_boundary_minimum():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

        def decode(self, ids, skip_special_tokens=True):
            del ids, skip_special_tokens
            return "。"

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    model.min_new_speak_tokens_before_chunk_boundary = 4
    model.max_new_speak_tokens_per_chunk = 64
    vocab_size = 151723
    alternative = 1234
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, alternative] = 20.0
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[151706, 200, 201, 202]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[alternative]]


def test_minicpmo_stage0_native_sampler_does_not_rewrite_model_punctuation():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        text = {
            108386: "你",
            104256: "好",
            3837: "，",
            100644: "今",
            99172: "天",
            100281: "想",
            27442: "聊",
            99217: "什",
            1773: "。",
            99218: "么",
        }

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

        def encode(self, text, add_special_tokens=False):
            del add_special_tokens
            return {"。": [1773]}.get(text, [])

        def decode(self, ids, skip_special_tokens=True):
            del skip_special_tokens
            return "".join(self.text.get(int(token_id), "") for token_id in ids)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    model.min_new_speak_tokens_before_chunk_boundary = 4
    model.max_new_speak_tokens_per_chunk = 64
    vocab_size = 151723
    period = 1773
    continuation = 99218
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, period] = 30.0
    logits[0, continuation] = 20.0
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[151706, 108386, 104256, 3837, 100644, 99172, 100281, 27442, 99217]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[period]]


def test_minicpmo_stage0_native_sampler_preserves_model_chunk_eos_decision():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    logits = torch.full((1, 151723), -100.0)
    logits[0, 151718] = 30.0
    logits[0, 1234] = 20.0
    sampling_metadata = SimpleNamespace(
        all_greedy=True,
        all_random=False,
        temperature=torch.tensor([0.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[151706, 200, 201]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[151718]]


def test_minicpmo_stage0_native_sampler_preserves_early_model_turn_eos_decision():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    model.min_new_speak_tokens_before_chunk_boundary = 8
    logits = torch.full((1, 151723), -100.0)
    logits[0, 151717] = 30.0
    logits[0, 1234] = 20.0
    sampling_metadata = SimpleNamespace(
        all_greedy=True,
        all_random=False,
        temperature=torch.tensor([0.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[151706, 200, 201]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[151717]]


def test_minicpmo_stage0_native_sampler_char_cap_does_not_override_special_token():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = [151706, 151717]

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

        def decode(self, ids, skip_special_tokens=True):
            del skip_special_tokens
            return "已经超过二十八个字符的当前语音分段文本"

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    logits = torch.full((1, 151723), -100.0)
    logits[0, 151717] = 30.0
    sampling_metadata = SimpleNamespace(
        all_greedy=True,
        all_random=False,
        temperature=torch.tensor([0.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[151706, 200, 201]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[151717]]


def test_minicpmo_stage0_native_sampler_preserves_model_chunk_eos():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    model.min_new_speak_tokens_before_chunk_boundary = 8
    model.max_new_speak_tokens_per_chunk = 64
    vocab_size = 151723
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, 151718] = 30.0
    logits[0, 1234] = 20.0
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[151706, 200, 201, 202]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[151718]]


def test_minicpmo_stage0_native_sampler_keeps_hard_chunk_cap():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import (
        MiniCPMO45DuplexPolicy,
    )
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

        def decode(self, ids, skip_special_tokens=True):
            del ids, skip_special_tokens
            return "没有自然边界"

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    vocab_size = 151723
    alternative = 1234
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, alternative] = 20.0
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[200] * (MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK - 1)],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK == 20
    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[151718]]


def test_minicpmo_stage0_native_sampler_cuts_before_request_length_cap():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

        def decode(self, ids, skip_special_tokens=True):
            del ids, skip_special_tokens
            return "没有自然边界"

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    model.max_new_speak_tokens_per_chunk = 64
    model._minicpmo45_duplex_row_max_tokens = {0: 20}
    vocab_size = 151723
    alternative = 1234
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, alternative] = 20.0
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[151706] + [200] * 18],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[151718]]


def test_minicpmo_stage0_native_sampler_cuts_on_decoded_text_length():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import (
        MiniCPMO45DuplexPolicy,
    )
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = [151706, 151718, 151717, 151705]

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

        def decode(self, ids, skip_special_tokens=True):
            token_text = {
                200: "一二三四五六七八九十",
                201: "十一十二十三十四十五",
                202: "十六十七十八十九",
            }
            special = set(self.all_special_ids) if skip_special_tokens else set()
            return "".join(token_text.get(int(token_id), "") for token_id in ids if int(token_id) not in special)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    vocab_size = 151723
    candidate = 202
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, candidate] = 20.0
    sampling_metadata = SimpleNamespace(
        all_greedy=True,
        all_random=False,
        temperature=torch.tensor([0.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[151706, 200, 201]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert MiniCPMO45DuplexPolicy.DEFAULT_MAX_SPEAK_CHARS_PER_CHUNK == 28
    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[151718]]


def test_minicpmo_stage0_native_sampler_ignores_pending_placeholders():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    vocab_size = 151723
    newline = 198
    alternative = 1234
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, newline] = 20.0
    logits[0, alternative] = 19.5
    sampling_metadata = SimpleNamespace(
        all_greedy=False,
        all_random=True,
        temperature=torch.tensor([1.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[-1, -1, -1]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[newline]]


def test_minicpmo_stage0_native_sampler_converts_mid_turn_listen_to_tts_bos():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

    state = SimpleNamespace(current_turn_ended=False)
    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    model._minicpmo45_duplex_row_sessions = {0: ("sid-native", 0)}
    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={("sid-native", 0): state})
    vocab_size = 151723
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, 151705] = 30.0
    sampling_metadata = SimpleNamespace(
        all_greedy=True,
        all_random=False,
        temperature=torch.tensor([0.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[151703]]
    assert state.current_turn_ended is False


def test_minicpmo_stage0_native_sampler_forced_listen_yields_floor():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    class _Tokenizer:
        eos_token_id = 151705
        unk_token_id = -1
        bad_token_ids = []
        all_special_ids = []

        def convert_tokens_to_ids(self, token):
            return {
                "<unit>": 151683,
                "</unit>": 151684,
                "<|listen|>": 151705,
                "<|speak|>": 151706,
                "<|tts_bos|>": 151703,
                "<|tts_eos|>": 151704,
                "<|tts_pad|>": 151722,
                "<|chunk_eos|>": 151718,
                "<|chunk_tts_eos|>": 151721,
                "<|turn_eos|>": 151717,
            }.get(token, -1)

    state = SimpleNamespace(current_turn_ended=False)
    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
    model._minicpmo45_duplex_row_sessions = {0: ("sid-native", 0)}
    model._minicpmo45_duplex_row_payloads = {0: {"force_listen": True}}
    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={("sid-native", 0): state})
    vocab_size = 151723
    logits = torch.full((1, vocab_size), -100.0)
    logits[0, 151705] = 30.0
    sampling_metadata = SimpleNamespace(
        all_greedy=True,
        all_random=False,
        temperature=torch.tensor([0.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[[]],
    )

    sampled = model.sample(logits, sampling_metadata)

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[151705]]
    assert state.current_turn_ended is True


def test_minicpmo_stage0_native_sampler_uses_runner_duplex_rows():
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    metadata = SimpleNamespace(
        prompt_token_ids=torch.tensor([[1, 2, 3]]),
    )

    rows = model._minicpmo45_native_duplex_prompt_rows(
        metadata,
        unit_id=151683,
        batch_size=1,
        duplex_rows=[0],
    )

    assert rows == [0]


def test_minicpmo_stage0_session_context_includes_resolved_ref_audio():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        MiniCPMO45Stage0DuplexRuntime,
        _MiniCPMO45Stage0SessionState,
    )

    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
    runtime.unit_token_id = 151683
    runtime.processor = SimpleNamespace()
    runtime.stage_model = SimpleNamespace()
    runtime.thinker = SimpleNamespace()
    runtime.device = "cpu"
    token_map = {
        "<|im_start|>system\nUse speech.\n<|audio_start|>": [1, 2, 3],
        "<|audio_end|><|im_end|>": [4, 5],
    }
    runtime._stage_runtime_ready = lambda: True
    runtime._require_special_token_ids = lambda: None
    runtime._decode_ref_audio_from_session_config = lambda _config: np.array([0.1, -0.1], dtype=np.float32)
    runtime._encode_text = lambda text: token_map[text]
    runtime._embed_token = lambda token_id: torch.full((1, 2), float(token_id))
    runtime._stage_ref_audio_embeddings = lambda ref_audio, state=None: torch.tensor([[10.0, 11.0], [12.0, 13.0]])

    state = _MiniCPMO45Stage0SessionState(session_id="sid-ref")
    runtime._prepare_session_context(state, {"instructions": "Use speech.", "extra_body": {"ref_audio_data": "x"}})

    assert state.context_token_ids == [1, 2, 3, 151683, 151683, 4, 5]
    assert len(state.context_embeds) == 6


# ---------------------------------------------------------------------------
# Cross-platform wiring: every AR model runner must reach the duplex hook.
#
# GPUARModelRunner and NPUARModelRunner are siblings -- neither inherits the
# other's _sample -- so the NPU runner once copied GPU _sample without the
# prepare_duplex_sampling call. Parse source instead of importing NPUARModelRunner:
# that class needs vllm_ascend, which CPU test images do not have.
# ---------------------------------------------------------------------------

_REQUIRED_DUPLEX_CALLS = {
    "__init__": "_init_duplex_sampling_state",
    "load_model": "_resolve_duplex_sampling_hook",
    "_update_states": "_update_duplex_sampling_states",
    "_sample": "_apply_duplex_sampling",
}

_MODEL_SAMPLER_CALL = "call_model_sampler"


def _repo_root():
    from pathlib import Path

    return Path(__file__).parents[2]


def _ar_runner_sources():
    root = _repo_root()
    paths = [root / "vllm_omni/worker/gpu_ar_model_runner.py"]
    paths.extend(sorted(root.glob("vllm_omni/platforms/*/worker/*_ar_model_runner.py")))
    return paths


def _ar_runner_classdefs():
    root = _repo_root()
    found = {}
    for path in _ar_runner_sources():
        module = ast.parse(path.read_text())
        for node in module.body:
            if isinstance(node, ast.ClassDef) and node.name.endswith("ARModelRunner"):
                found[node.name] = (node, str(path.relative_to(root)))
    assert found, "no AR model runners discovered; the glob in _ar_runner_sources() is stale"
    return found


def _method_calls(method):
    calls: dict[str, int] = {}
    for node in ast.walk(method):
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
            calls.setdefault(node.func.attr, node.lineno)
    return calls


def _wires_duplex_sampling_directly(classdef) -> bool:
    bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)}
    if "DuplexSamplingRunnerMixin" not in bases:
        return False
    methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)}
    return all(
        method_name in methods and required_call in _method_calls(methods[method_name])
        for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items()
    )


def test_every_ar_runner_reaches_the_duplex_sampling_hook():
    classdefs = _ar_runner_classdefs()
    wired = {name for name, (classdef, _path) in classdefs.items() if _wires_duplex_sampling_directly(classdef)}

    for class_name, (classdef, path) in sorted(classdefs.items()):
        if class_name in wired:
            continue
        base_names = {base.id for base in classdef.bases if isinstance(base, ast.Name)}
        assert base_names & wired, (
            f"{path}::{class_name} neither wires the duplex sampling hook nor derives from a "
            f"runner that does (wired runners: {sorted(wired)}).  Inherit DuplexSamplingRunnerMixin "
            f"and call {sorted(_REQUIRED_DUPLEX_CALLS.values())} from "
            f"{sorted(_REQUIRED_DUPLEX_CALLS)} respectively."
        )


@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"])
def test_ar_runner_calls_every_duplex_sampling_hook_site(class_name: str):
    classdefs = _ar_runner_classdefs()
    assert class_name in classdefs, f"{class_name} not found by _ar_runner_sources()"
    classdef, path = classdefs[class_name]

    bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)}
    assert "DuplexSamplingRunnerMixin" in bases, (
        f"{path}::{class_name} must inherit DuplexSamplingRunnerMixin so the duplex "
        f"sampling hook is wired the same way on every platform"
    )

    methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)}
    for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items():
        method = methods.get(method_name)
        assert method is not None, f"{path}::{class_name}.{method_name} is missing; it must call {required_call}"
        calls = _method_calls(method)
        assert required_call in calls, f"{path}::{class_name}.{method_name} must call self.{required_call}"


@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"])
def test_ar_runner_applies_duplex_sampling_before_the_model_sampler(class_name: str):
    classdef, path = _ar_runner_classdefs()[class_name]
    sample = next(node for node in classdef.body if isinstance(node, ast.FunctionDef) and node.name == "_sample")

    apply_lineno = _method_calls(sample).get("_apply_duplex_sampling")
    sampler_lineno = next(
        (
            node.lineno
            for node in ast.walk(sample)
            if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == _MODEL_SAMPLER_CALL
        ),
        None,
    )
    assert apply_lineno is not None, f"{path}::{class_name}._sample must call self._apply_duplex_sampling"
    assert sampler_lineno is not None, (
        f"{path}::{class_name}._sample no longer calls {_MODEL_SAMPLER_CALL}(); update _MODEL_SAMPLER_CALL "
        f"so the ordering check keeps testing something"
    )
    assert apply_lineno < sampler_lineno, (
        f"{path}::{class_name}._sample calls self._apply_duplex_sampling at line {apply_lineno}, after "
        f"{_MODEL_SAMPLER_CALL}() at line {sampler_lineno}.  The hook masks the logits and publishes the "
        f"row -> session map the sampler reads, so running it afterwards is a silent no-op."
    )


class _NativeDuplexTokenizer:
    eos_token_id = 151705
    unk_token_id = -1
    bad_token_ids: list[int] = []
    all_special_ids: list[int] = []

    def convert_tokens_to_ids(self, token):
        return {
            "<unit>": 151683,
            "</unit>": 151684,
            "<|listen|>": 151705,
            "<|speak|>": 151706,
            "<|tts_bos|>": 151703,
            "<|tts_eos|>": 151704,
            "<|tts_pad|>": 151722,
            "<|chunk_eos|>": 151718,
            "<|chunk_tts_eos|>": 151721,
            "<|turn_eos|>": 151717,
        }.get(token, -1)


def _native_duplex_model_with_state():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
        _MiniCPMO45Stage0SessionState,
    )
    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
        MiniCPMO45OmniForConditionalGeneration,
    )

    session_key = "sid-lookahead"
    state = _MiniCPMO45Stage0SessionState(session_id=session_key, current_turn_ended=False)
    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
    model.model_stage = "llm"
    model.thinker = SimpleNamespace(get_tokenizer=lambda: _NativeDuplexTokenizer())
    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state})
    model._minicpmo45_duplex_row_sessions = {0: session_key}
    model._minicpmo45_duplex_row_payloads = {0: {"is_speech": True}}
    model._minicpmo45_active_duplex_rows = [0]
    return model, state


def _native_duplex_sampling_metadata(output_token_ids: list[int]):
    return SimpleNamespace(
        all_greedy=True,
        all_random=False,
        temperature=torch.tensor([0.0]),
        top_k=torch.tensor([1]),
        top_p=torch.tensor([1.0]),
        generators={},
        prompt_token_ids=torch.tensor([[151683] * 16]),
        output_token_ids=[output_token_ids],
    )


@pytest.mark.parametrize("terminator", [151718, 151721, 151705])
def test_minicpmo_stage0_async_lookahead_frame_after_chunk_terminator_is_frozen(terminator):
    """The frame async scheduling runs after a sampled chunk terminator is
    discarded by the scheduler; it must neither decide anything nor touch the
    session state that the next append re-injects."""
    model, state = _native_duplex_model_with_state()
    state.pending_terminator_token = terminator
    state.last_terminator_token = terminator
    state.generated_tokens = [200, 201]
    logits = torch.full((1, 151723), -100.0)
    logits[0, 1234] = 30.0  # what the stale frame would otherwise pick

    sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201, terminator, -1]))

    assert sampled is not None
    assert sampled.sampled_token_ids.tolist() == [[terminator]]
    assert state.pending_terminator_token == terminator
    assert state.last_terminator_token == terminator
    assert state.generated_tokens == [200, 201]
    assert state.current_turn_ended is False


def test_minicpmo_stage0_turn_eos_is_forwarded_not_pending():
    """<|turn_eos|> is fed like text in the official unit loop; only the
    chunk terminator sampled afterwards is re-injected on the next append."""
    model, state = _native_duplex_model_with_state()
    logits = torch.full((1, 151723), -100.0)
    logits[0, 151717] = 30.0

    sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201]))

    assert sampled.sampled_token_ids.tolist() == [[151717]]
    assert state.pending_terminator_token is None
    assert state.last_terminator_token == 151717
    assert state.current_turn_ended is True

    # The frame that consumes <|turn_eos|> keeps sampling; the chunk
    # terminator it picks becomes the pending token for the next append.
    logits = torch.full((1, 151723), -100.0)
    logits[0, 151718] = 30.0
    sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201, 151717]))

    assert sampled.sampled_token_ids.tolist() == [[151718]]
    assert state.pending_terminator_token == 151718
    assert state.current_turn_ended is True


def test_minicpmo_stage0_stop_tokens_exclude_turn_eos():
    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import _stage0_stop_token_ids

    stop_ids = _stage0_stop_token_ids(_NativeDuplexTokenizer())

    assert stop_ids == [151718, 151721, 151705]