vllm-omni/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

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

import threading
import uuid
from collections import deque
from types import MethodType, SimpleNamespace
from unittest.mock import Mock, call, patch

import pytest
import torch
from pytest_mock import MockerFixture
from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
from vllm.v1.request import Request, RequestStatus

from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt
from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
    OmniChunkTransferAdapter,
)
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec

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


@pytest.fixture
def shm_sender(build_adapter):
    adapter, _ = build_adapter(stage_id=0)
    connector = SharedMemoryConnector({"stage_id": 0})
    adapter.connector = connector
    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
        codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
    )
    yield adapter, connector
    connector.close()


def test_shm_abort_cleans_only_owned_chunks_before_id_reuse(shm_sender):
    adapter, connector = shm_sender
    ext_id = f"abort_{uuid.uuid4().hex}"
    request = _req("old", RequestStatus.RUNNING, external_req_id=ext_id)
    sibling_key = f"{ext_id}_1_2_0_0"
    assert connector.put("0", "1", sibling_key, "sibling")[0]
    adapter.save_async(None, request)
    adapter._send_single_request(adapter._pending_save_reqs.popleft())

    adapter.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED, {request.request_id: request})
    # Queue a replacement before cleanup runs. The same save loop must reclaim
    # the old generation before its key can be written by the new generation.
    replacement = _req("new", RequestStatus.RUNNING, external_req_id=ext_id)
    adapter.save_async(None, replacement)
    adapter._send_single_request(adapter._pending_save_reqs.popleft())
    assert connector.get("0", "1", f"{ext_id}_0_0") is None
    adapter._send_single_request(adapter._pending_save_reqs.popleft())
    assert connector.get("0", "1", f"{ext_id}_0_0") is not None
    assert connector.get("0", "1", sibling_key)[0] == "sibling"


def test_shm_abort_reclaims_inflight_put_without_blocking(shm_sender, monkeypatch):
    adapter, connector = shm_sender
    request = _req("inflight", RequestStatus.RUNNING, external_req_id=f"inflight_{uuid.uuid4().hex}")
    put_started = threading.Event()
    release_put = threading.Event()
    real_put = connector.put

    def blocking_put(**kwargs):
        put_started.set()
        assert release_put.wait(timeout=5)
        return real_put(**kwargs)

    monkeypatch.setattr(connector, "put", blocking_put)
    adapter.save_async(None, request)
    task = adapter._pending_save_reqs.popleft()
    sender = threading.Thread(target=adapter._send_single_request, args=(task,))
    sender.start()
    try:
        assert put_started.wait(timeout=5)
        adapter.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED, {request.request_id: request})
        assert sender.is_alive()
    finally:
        release_put.set()
        sender.join(timeout=5)
    assert not sender.is_alive()
    while adapter._pending_save_reqs:
        adapter._send_single_request(adapter._pending_save_reqs.popleft())
    assert not connector._pending_keys
    assert connector.get("0", "1", f"{request.external_req_id}_0_0") is None
    assert not adapter._sender_tokens
    assert not adapter.code_prompt_token_ids


def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender):
    adapter, connector = shm_sender
    ext_id = f"release_{uuid.uuid4().hex}"
    request = _req("done", RequestStatus.RUNNING, external_req_id=ext_id)
    sibling_key = f"{ext_id}_1_0_0"
    assert connector.put("0", "1", sibling_key, "sibling")[0]
    for _ in range(3):
        adapter.save_async(None, request)
        adapter._send_single_request(adapter._pending_save_reqs.popleft())
    # The consumer read only the first chunk before its own request ended.
    assert connector.get("0", "1", f"{ext_id}_0_0") is not None

    adapter.release_shm_resources(ext_id)
    # The release is queued behind sends; nothing is unlinked on the caller.
    assert f"{ext_id}_0_1" in connector._pending_keys
    while adapter._pending_save_reqs:
        adapter._send_single_request(adapter._pending_save_reqs.popleft())

    assert adapter._reclaimed_shm_total == 2
    assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys)
    assert connector.get("0", "1", f"{ext_id}_0_1") is None
    assert connector.get("0", "1", sibling_key)[0] == "sibling"


@pytest.mark.parametrize("fail_terminal", [False, True])
def test_shm_terminal_keeps_payload_only_after_success(shm_sender, monkeypatch, fail_terminal):
    adapter, connector = shm_sender
    request = _req("terminal", RequestStatus.FINISHED_STOPPED, external_req_id=f"terminal_{uuid.uuid4().hex}")
    if fail_terminal:
        real_put = connector.put

        def failed_put(**kwargs):
            real_put(**kwargs)
            return False, 0, None

        monkeypatch.setattr(connector, "put", failed_put)
    adapter.save_async(None, request)
    adapter.cleanup_sender(request.external_req_id)
    while adapter._pending_save_reqs:
        adapter._send_single_request(adapter._pending_save_reqs.popleft())
    result = connector.get("0", "1", f"{request.external_req_id}_0_0")
    assert (result is None) == fail_terminal
    assert not adapter._sender_tokens


@pytest.mark.parametrize("stage_id", [0, 1])
def test_receiver_cleanup_does_not_accumulate_tombstones(build_adapter, stage_id):
    adapter, _ = build_adapter(stage_id=stage_id)
    for index in range(256):
        adapter.cleanup_receiver(f"completed-{index}")
    assert not getattr(adapter, "_cancelled_load_reqs", ())
    assert not adapter._registered_load_entries


def test_recv_loop_does_not_recreate_mapping_for_cancelled_entry(build_adapter):
    adapter, connector = build_adapter(stage_id=1)
    request = _req("cancelled", RequestStatus.RUNNING, external_req_id="external")
    adapter.load_async(request)
    adapter.cleanup_receiver(request.request_id)
    # Stop after the loop has drained the stale queue entry.
    adapter._recv_cond.wait = lambda **kwargs: adapter.stop_event.set()
    adapter.recv_loop()
    connector.get.assert_not_called()
    assert not adapter.request_ids_mapping
    assert not adapter._pending_load_reqs


def test_idle_save_loop_reaps_consumed_shm(shm_sender):
    adapter, connector = shm_sender
    key = f"idle_{uuid.uuid4().hex}"
    assert connector.put("0", "1", key, "payload")[0]
    receiver = SharedMemoryConnector({})
    try:
        assert receiver.get("0", "1", key)[0] == "payload"
        adapter._save_cond.wait = lambda **kwargs: adapter.stop_event.set()
        adapter.save_loop()
        assert not connector._pending_keys
    finally:
        receiver.close()


class DummyWaitingQueue(list):
    def prepend_requests(self, requests):
        self[:0] = list(requests)

    def add_request(self, request):
        self.append(request)

    def remove_requests(self, requests):
        remove = set(requests)
        self[:] = [request for request in self if request not in remove]


def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None):
    request = Mock(
        client_index=0,
        request_id=req_id,
        external_req_id=external_req_id or req_id,
        status=status,
        prompt_token_ids=[],
        num_prompt_tokens=0,
        num_computed_tokens=0,
        num_output_placeholders=0,
        prefill_stats=None,
        additional_information=None,
        resumable=False,
    )
    request.is_finished = lambda: RequestStatus.is_finished(request.status)
    return request


def _streaming_request(mocker: MockerFixture, num_computed_tokens: int) -> SimpleNamespace:
    prompt_token_ids = [0] * num_computed_tokens
    return SimpleNamespace(
        _all_token_ids=prompt_token_ids.copy(),
        _output_token_ids=[],
        prompt_token_ids=prompt_token_ids,
        num_computed_tokens=num_computed_tokens,
        num_prompt_tokens=num_computed_tokens,
        num_output_placeholders=0,
        update_block_hashes=mocker.Mock(),
    )


def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture) -> None:
    request = SimpleNamespace(
        _all_token_ids=[0, 0, 7, 8],
        _output_token_ids=[7, 8],
        prompt_token_ids=[0, 0],
        num_computed_tokens=4,
        num_prompt_tokens=2,
        update_block_hashes=mocker.Mock(),
    )
    payload = {
        "ids": {"prompt": [1, 2, 3]},
        "meta": {
            "replace_streaming_prompt": True,
            "next_stage_prompt_len": 7,
        },
    }

    construct_next_stage_streaming_input_prompt(payload, request)

    assert request.prompt_token_ids == [0] * 7
    assert request._all_token_ids == [0] * 7
    assert request._output_token_ids == []
    assert request.num_computed_tokens == 0
    assert request.num_prompt_tokens == 7
    request.update_block_hashes.assert_called_once_with()


def test_turn_start_replacement_ignores_accumulated_prompt_capacity(mocker: MockerFixture) -> None:
    request = _streaming_request(mocker, num_computed_tokens=4064)
    payload: dict = {
        "ids": {"prompt": [1]},
        "meta": {
            "replace_streaming_prompt": True,
            "next_stage_prompt_len": 10,
            "next_stage_generation_tokens": 26,
        },
    }

    replaced = construct_next_stage_streaming_input_prompt(
        payload,
        request,
        max_model_len=4096,
        previous_condition_len=10,
        previous_condition_seq=7,
        condition_seq=8,
        recompute_previous_chunks=1,
    )

    assert replaced is True
    assert request.prompt_token_ids == [0] * 10
    assert payload["meta"]["streaming_prompt_recompute"] is False
    assert "streaming_prompt_previous_codes" not in payload["ids"]


@pytest.mark.parametrize("reserve", [True, 0, -1, 1.5, "26", None])
def test_streaming_prompt_rejects_invalid_generation_reserve(
    mocker: MockerFixture,
    reserve: object,
) -> None:
    request = _streaming_request(mocker, num_computed_tokens=10)
    payload = {
        "ids": {"prompt": [1]},
        "meta": {
            "next_stage_prompt_len": 10,
            "next_stage_generation_tokens": reserve,
        },
    }

    with pytest.raises(ValueError, match="next_stage_generation_tokens must be a positive integer"):
        construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)


def test_capacity_managed_streaming_window_at_limit_still_appends(mocker: MockerFixture) -> None:
    request = _streaming_request(mocker, num_computed_tokens=4060)
    payload = {
        "ids": {"prompt": [1]},
        "meta": {
            "next_stage_prompt_len": 10,
            "next_stage_generation_tokens": 26,
        },
    }

    replaced = construct_next_stage_streaming_input_prompt(
        payload,
        request,
        max_model_len=4096,
        previous_condition_len=10,
        previous_condition_seq=0,
        condition_seq=1,
        recompute_previous_chunks=1,
        recompute_on_capacity=True,
    )

    assert replaced is False
    assert request.num_computed_tokens == 4060
    assert request.num_prompt_tokens == 4070


def test_capacity_managed_streaming_prompt_appends_from_declared_length_without_ids_prompt(
    mocker: MockerFixture,
) -> None:
    request = _streaming_request(mocker, num_computed_tokens=20)
    payload = {
        "ids": {"tts": [1, 2, 3]},
        "hidden_states": {"tts": [[0.1], [0.2], [0.3]]},
        "meta": {
            "next_stage_prompt_len": 4,
            "next_stage_generation_tokens": 26,
        },
    }

    replaced = construct_next_stage_streaming_input_prompt(
        payload,
        request,
        max_model_len=4096,
        previous_condition_len=4,
        previous_condition_seq=0,
        condition_seq=1,
        recompute_previous_chunks=1,
        recompute_on_capacity=True,
    )

    assert replaced is False
    assert request.num_computed_tokens == 20
    assert request.num_prompt_tokens == 24
    assert request.prompt_token_ids == [0] * 24
    request.update_block_hashes.assert_called_once_with()


def test_streaming_window_builds_one_chunk_recompute_recipe(mocker: MockerFixture) -> None:
    previous_condition_len = 10
    previous_codes = list(range(25))
    request = SimpleNamespace(
        _all_token_ids=[0] * 4039 + previous_codes + [999],
        _output_token_ids=previous_codes + [999],
        prompt_token_ids=[0] * 4039,
        num_computed_tokens=4065,
        num_prompt_tokens=4039,
        num_output_placeholders=1,
        update_block_hashes=mocker.Mock(),
    )
    payload: dict = {
        "ids": {"prompt": [1]},
        "meta": {
            "next_stage_prompt_len": 10,
            "next_stage_generation_tokens": 26,
        },
    }

    replaced = construct_next_stage_streaming_input_prompt(
        payload,
        request,
        max_model_len=4096,
        previous_condition_len=previous_condition_len,
        previous_condition_seq=7,
        condition_seq=8,
        recompute_previous_chunks=1,
    )

    assert replaced is True
    assert request.prompt_token_ids == [0] * 45
    assert request.num_prompt_tokens == 45
    assert request.num_computed_tokens == 0
    assert payload["ids"]["streaming_prompt_previous_codes"] == previous_codes
    assert payload["meta"] == {
        "next_stage_prompt_len": 10,
        "next_stage_generation_tokens": 26,
        "streaming_prompt_recompute": True,
        "streaming_condition_seq": 8,
    }
    request.update_block_hashes.assert_called_once_with()


def test_capacity_policy_turn_start_replacement_does_not_require_generation_reserve(
    mocker: MockerFixture,
) -> None:
    request = _streaming_request(mocker, num_computed_tokens=4064)
    payload = {
        "ids": {"prompt": [1]},
        "meta": {
            "replace_streaming_prompt": True,
            "next_stage_prompt_len": 10,
        },
    }

    replaced = construct_next_stage_streaming_input_prompt(
        payload,
        request,
        max_model_len=4096,
        previous_condition_len=10,
        previous_condition_seq=7,
        condition_seq=8,
        recompute_previous_chunks=1,
        recompute_on_capacity=True,
    )

    assert replaced is True
    assert request.num_computed_tokens == 0
    assert request.num_prompt_tokens == 10


def test_streaming_window_appends_until_capacity_then_recomputes(mocker: MockerFixture) -> None:
    request = SimpleNamespace(
        _all_token_ids=[0] * 4000 + [101, 102],
        _output_token_ids=[101, 102],
        prompt_token_ids=[0] * 4000,
        num_computed_tokens=4002,
        num_prompt_tokens=4000,
        num_output_placeholders=0,
        update_block_hashes=mocker.Mock(),
    )
    second: dict = {
        "ids": {"prompt": [1]},
        "meta": {
            "next_stage_prompt_len": 12,
            "next_stage_generation_tokens": 26,
        },
    }

    assert not construct_next_stage_streaming_input_prompt(
        second,
        request,
        max_model_len=4096,
        previous_condition_len=10,
        previous_condition_seq=0,
        condition_seq=1,
        recompute_previous_chunks=1,
        recompute_on_capacity=True,
    )
    assert request.num_computed_tokens == 4002
    assert request.num_prompt_tokens == 4014
    assert "streaming_prompt_previous_codes" not in second["ids"]
    assert second["meta"]["streaming_prompt_recompute"] is False

    request._all_token_ids.extend([201, 202, 203])
    request._output_token_ids.extend([201, 202, 203])
    request.num_computed_tokens = 4017
    third: dict = {
        "ids": {"prompt": [2]},
        "meta": {
            "next_stage_prompt_len": 60,
            "next_stage_generation_tokens": 26,
        },
    }

    assert construct_next_stage_streaming_input_prompt(
        third,
        request,
        max_model_len=4096,
        previous_condition_len=12,
        previous_condition_seq=1,
        condition_seq=2,
        recompute_previous_chunks=1,
        recompute_on_capacity=True,
    )
    assert request.num_prompt_tokens == 75
    assert request.num_computed_tokens == 0
    assert third["ids"]["streaming_prompt_previous_codes"] == [201, 202, 203]
    assert third["meta"]["streaming_condition_seq"] == 2

    request._all_token_ids.extend([301])
    request._output_token_ids.extend([301])
    request.num_computed_tokens = 76
    fourth: dict = {
        "ids": {"prompt": [3]},
        "meta": {
            "next_stage_prompt_len": 8,
            "next_stage_generation_tokens": 26,
            "streaming_condition_seq": 3,
        },
    }

    assert not construct_next_stage_streaming_input_prompt(
        fourth,
        request,
        max_model_len=4096,
        previous_condition_len=60,
        previous_condition_seq=2,
        condition_seq=3,
        recompute_previous_chunks=1,
        recompute_on_capacity=True,
    )
    assert request.num_computed_tokens == 76
    assert request.num_prompt_tokens == 84
    assert fourth["meta"]["streaming_prompt_recompute"] is False


def test_capacity_rollover_requires_explicit_window_contract(mocker: MockerFixture) -> None:
    request = _streaming_request(mocker, num_computed_tokens=4064)
    payload = {
        "ids": {"prompt": [1]},
        "meta": {
            "next_stage_prompt_len": 10,
            "next_stage_generation_tokens": 26,
        },
    }

    with pytest.raises(ValueError, match="requires window_size=1"):
        construct_next_stage_streaming_input_prompt(
            payload,
            request,
            max_model_len=4096,
            previous_condition_len=10,
            previous_condition_seq=0,
            condition_seq=1,
        )


@pytest.mark.parametrize("next_stage_prompt_len", [0, -1])
def test_capacity_managed_streaming_prompt_rejects_nonpositive_length(
    mocker: MockerFixture, next_stage_prompt_len: int
) -> None:
    request = _streaming_request(mocker, num_computed_tokens=0)
    payload = {
        "ids": {"prompt": [1]},
        "meta": {
            "next_stage_prompt_len": next_stage_prompt_len,
            "next_stage_generation_tokens": 26,
        },
    }

    with pytest.raises(ValueError, match="positive next_stage_prompt_len"):
        construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)


def test_capacity_managed_streaming_prompt_rejects_fresh_overflow(mocker: MockerFixture) -> None:
    request = _streaming_request(mocker, num_computed_tokens=0)
    payload = {
        "ids": {"prompt": [1]},
        "meta": {
            "next_stage_prompt_len": 4071,
            "next_stage_generation_tokens": 26,
        },
    }

    with pytest.raises(ValueError, match="exceeds max_model_len"):
        construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)


@pytest.fixture
def build_adapter(monkeypatch, mocker: MockerFixture):
    def _build(
        *,
        stage_id: int = 1,
        model_mode: str = "ar",
        max_num_seqs: int = 2,
        max_model_len: int = 0,
        tts_max_model_len: int = 0,
        tts_attention_type: str = "full_attention",
        flat_tts_config: bool = False,
        active_stream_window: int = 0,
        connector_extra: dict | None = None,
    ):
        connector = mocker.MagicMock()
        connector.stage_id = stage_id
        connector.config = {"extra": connector_extra or {}}
        connector.get.return_value = None
        connector.put.return_value = (True, 1, {})

        def _fake_base_init(self, config):
            self.config = config
            self._pending_load_reqs = deque()
            self._finished_load_reqs = set()
            self._cancelled_load_reqs = set()
            self._pending_save_reqs = deque()
            self._finished_save_reqs = set()
            self._send_failures = {}
            self._send_failure_lock = threading.Lock()
            self._receive_failures = {}
            self._receive_failure_lock = threading.Lock()
            self.stop_event = threading.Event()
            self._recv_cond = threading.Condition()
            self._save_cond = threading.Condition()

        monkeypatch.setattr(OmniTransferAdapterBase, "__init__", _fake_base_init)
        monkeypatch.setattr(
            OmniChunkTransferAdapter,
            "create_connector",
            classmethod(lambda cls, _model_config: connector),
        )

        hf_config = SimpleNamespace(
            model_type="conditional_chattts",
            max_position_embeddings=tts_max_model_len,
            attention_type=tts_attention_type,
        )
        if not flat_tts_config:
            hf_config = SimpleNamespace(
                tts_config=SimpleNamespace(
                    model_type="minicpmtts",
                    max_position_embeddings=tts_max_model_len,
                    attention_type=tts_attention_type,
                )
            )
        model_config = SimpleNamespace(
            worker_type=model_mode,
            max_num_seqs=max_num_seqs,
            max_model_len=max_model_len,
            hf_config=hf_config,
            hf_config_name="tts_config" if stage_id == 1 else None,
            active_stream_window=active_stream_window,
            stage_connector_config={
                "name": "SharedMemoryConnector",
                "extra": connector_extra or {},
            },
        )
        scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs)
        adapter = OmniChunkTransferAdapter(
            SimpleNamespace(model_config=model_config, scheduler_config=scheduler_config)
        )
        return adapter, connector

    return _build


def _dequeue_load_entry(adapter, request):
    """Register and take the entry as the background recv loop would."""
    adapter.load_async(request)
    entry = adapter._registered_load_entries[request.request_id]
    try:
        adapter._pending_load_reqs.remove(entry)
    except ValueError:
        pass
    return entry


@pytest.mark.parametrize(
    ("attention_type", "expected_previous_chunks", "expected_capacity_trigger"),
    [("full_attention", 1, True), ("sliding_recompute", 1, False), ("other", 0, False)],
)
def test_talker_attention_policy_controls_streaming_recompute(
    build_adapter, attention_type, expected_previous_chunks, expected_capacity_trigger
):
    adapter, _ = build_adapter(tts_attention_type=attention_type)

    assert adapter._streaming_prompt_previous_chunks == expected_previous_chunks
    assert adapter._streaming_prompt_recompute_on_capacity is expected_capacity_trigger


@pytest.mark.parametrize("stage_id", [0, 2])
def test_talker_streaming_policy_is_scoped_to_tts_stage(build_adapter, stage_id: int) -> None:
    adapter, _ = build_adapter(stage_id=stage_id, tts_attention_type="full_attention")

    assert adapter._streaming_prompt_previous_chunks == 0
    assert adapter._streaming_prompt_recompute_on_capacity is False


@pytest.mark.parametrize(
    ("raw_cfg", "expected_name", "expected_extra"),
    [
        (None, "SharedMemoryConnector", {}),
        (SimpleNamespace(name="YuanrongConnector", extra={"k": "v"}), "YuanrongConnector", {"k": "v"}),
    ],
)
def test_create_connector_config_parsing(monkeypatch, raw_cfg, expected_name, expected_extra):
    captured = {}

    def _fake_create(spec):
        captured["spec"] = spec
        return "ok"

    monkeypatch.setattr(
        "vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter"
        ".OmniConnectorFactory.create_connector",
        _fake_create,
    )

    model_config = SimpleNamespace(stage_connector_config=raw_cfg) if raw_cfg is not None else SimpleNamespace()
    connector = OmniChunkTransferAdapter.create_connector(model_config)

    assert connector == "ok"
    assert isinstance(captured["spec"], ConnectorSpec)
    assert captured["spec"].name == expected_name
    assert captured["spec"].extra == {**expected_extra, "stage_id": 0}


def test_load_poll(build_adapter):
    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
    request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")

    adapter.load_async(request)
    payload: OmniPayload = {
        "codes": {"audio": [[1]]},
        "hidden_states": {"output": torch.tensor([[2.0]])},
        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
    }
    connector.get.return_value = (payload, 16)
    adapter._poll_single_request(_dequeue_load_entry(adapter, request))

    assert request.additional_information is None
    assert "req-1" in adapter._finished_load_reqs
    request.status = RequestStatus.WAITING_FOR_CHUNK
    adapter.process_pending_chunks(
        DummyWaitingQueue(),
        [request],
        scheduler_requests={request.request_id: request},
    )
    assert request.additional_information == payload
    assert adapter.get_req_chunk["req-1"] == 1
    assert "req-1" not in adapter._finished_load_reqs


def test_load_poll_uses_request_scoped_payload_sender_endpoint(build_adapter):
    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
    request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
    request.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
    connector.get.return_value = None

    adapter._poll_single_request(_dequeue_load_entry(adapter, request))

    connector.get.assert_called_once_with(
        "1",
        "2",
        "external-1_1_0",
        {"source_host": "10.0.0.1", "source_port": 50051},
    )


def test_load_poll_keeps_concurrent_payload_sender_endpoints_distinct(build_adapter):
    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
    first = _req("req-1", RequestStatus.WAITING)
    first.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 50051}
    second = _req("req-2", RequestStatus.WAITING)
    second.payload_sender_info = {"host": "10.0.0.2", "zmq_port": 51051}
    connector.get.return_value = None

    second_entry = _dequeue_load_entry(adapter, second)
    first_entry = _dequeue_load_entry(adapter, first)
    adapter._poll_single_request(second_entry)
    adapter._poll_single_request(first_entry)

    assert connector.get.call_args_list == [
        call(
            "1",
            "2",
            "req-2_1_0",
            {"source_host": "10.0.0.2", "source_port": 51051},
        ),
        call(
            "1",
            "2",
            "req-1_1_0",
            {"source_host": "10.0.0.1", "source_port": 50051},
        ),
    ]


def test_load_async_does_not_requeue_registered_inflight_request(build_adapter):
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    request = _req("req-inflight", RequestStatus.RUNNING, external_req_id="ext-inflight")

    entry = _dequeue_load_entry(adapter, request)

    # The recv loop owns the popped request until polling completes. A
    # scheduler-side duplicate registration must not append a second copy.
    adapter.load_async(request)
    assert adapter._pending_load_reqs == deque()

    connector.get.return_value = (
        {
            "ids": {"prompt": [1]},
            "meta": {"finished": True},
        },
        1,
    )
    assert adapter._poll_single_request(entry) is True

    # Successful polling releases the registration for the next chunk.
    adapter.load_async(request)
    assert [queued.request for queued in adapter._pending_load_reqs] == [request]


def test_cleanup_then_reregister_drops_stale_queued_entry(build_adapter):
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    request = _req("req-resume", RequestStatus.WAITING, external_req_id="ext-resume")

    adapter.load_async(request)
    stale_entry = adapter._pending_load_reqs[0]
    adapter.cleanup_receiver(request.request_id)
    adapter.load_async(request)
    fresh_entry = adapter._pending_load_reqs[-1]
    assert len(adapter._pending_load_reqs) == 2
    assert adapter._pending_load_reqs.popleft() is stale_entry
    assert adapter._pending_load_reqs.popleft() is fresh_entry

    connector.get.return_value = (
        {
            "ids": {"prompt": [1]},
            "meta": {"finished": True},
        },
        1,
    )
    assert adapter._poll_single_request(stale_entry) is True
    connector.get.assert_not_called()
    assert adapter._registered_load_entries[request.request_id] is fresh_entry

    assert adapter._poll_single_request(fresh_entry) is True
    connector.get.assert_called_once()
    assert request.request_id not in adapter._registered_load_entries
    assert adapter.get_req_chunk[request.request_id] == 1


def test_load_poll_ar_requeues_explicitly_replaced_running_prompt(build_adapter):
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    request = _req("req-replace", RequestStatus.RUNNING, external_req_id="external-replace")
    request.resumable = True
    request.prompt_token_ids = [0] * 4
    request._all_token_ids = [0] * 4
    request._output_token_ids = []
    request.num_prompt_tokens = 4
    request.num_computed_tokens = 4
    request.update_block_hashes = Mock()
    adapter.get_req_chunk[request.request_id] = 1
    adapter.requests_num_chunks_sent[request.external_req_id] = 4
    adapter.request_ids_mapping[request.request_id] = request.external_req_id
    connector.get.return_value = (
        {
            "ids": {"prompt": [1]},
            "meta": {
                "next_stage_prompt_len": 10,
                "replace_streaming_prompt": True,
                "finished": False,
            },
        },
        1,
    )

    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
    request.status = RequestStatus.WAITING_FOR_CHUNK
    running_queue = [request]
    waiting_queue = DummyWaitingQueue()
    adapter.process_pending_chunks(
        waiting_queue,
        running_queue,
        scheduler_requests={request.request_id: request},
    )
    assert request.num_computed_tokens == 0
    assert request.prompt_token_ids == [0] * 10
    assert request.request_id in adapter.replaced_streaming_prompt_ids
    assert running_queue == []
    assert waiting_queue == [request]
    assert request.status == RequestStatus.WAITING

    adapter.requests_num_chunks_sent[request.external_req_id] = 9
    adapter.postprocess_scheduler_output(
        SimpleNamespace(
            scheduled_new_reqs=[SimpleNamespace(req_id=request.request_id)],
            scheduled_cached_reqs=SimpleNamespace(req_ids=[]),
        )
    )
    assert request.request_id not in adapter.replaced_streaming_prompt_ids
    assert request.external_req_id not in adapter.requests_num_chunks_sent


@pytest.mark.parametrize("flat_tts_config", [False, True])
def test_load_poll_ar_recomputes_native_duplex_prompt_over_capacity(build_adapter, flat_tts_config):
    adapter, connector = build_adapter(
        stage_id=1,
        model_mode="ar",
        max_model_len=8192,
        tts_max_model_len=4096,
        tts_attention_type="full_attention",
        flat_tts_config=flat_tts_config,
    )
    request = _req("req-rollover", RequestStatus.RUNNING, external_req_id="external-rollover")
    request.resumable = True
    previous_codes = list(range(25))
    request.prompt_token_ids = [0] * 4039
    request._all_token_ids = [0] * 4039 + previous_codes
    request._output_token_ids = previous_codes.copy()
    request.num_prompt_tokens = 4039
    request.num_computed_tokens = 4064
    request.num_output_placeholders = 0
    request.update_block_hashes = Mock()
    adapter.get_req_chunk[request.request_id] = 1
    adapter.request_ids_mapping[request.request_id] = request.external_req_id
    adapter._streaming_condition_lengths[request.request_id] = 10
    adapter._streaming_condition_seqs[request.request_id] = 0
    connector.get.return_value = (
        {
            "ids": {"prompt": [1]},
            "meta": {
                "next_stage_prompt_len": 10,
                "next_stage_generation_tokens": 26,
                "finished": True,
            },
        },
        1,
    )

    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
    request.status = RequestStatus.WAITING_FOR_CHUNK
    running_queue = [request]
    waiting_queue = DummyWaitingQueue()
    adapter.process_pending_chunks(
        waiting_queue,
        running_queue,
        scheduler_requests={request.request_id: request},
    )

    assert request.num_computed_tokens == 0
    assert request.prompt_token_ids == [0] * 45
    assert request.additional_information["ids"]["streaming_prompt_previous_codes"] == previous_codes
    assert request.additional_information["meta"]["streaming_prompt_recompute"] is True
    assert request.resumable is False
    assert request.request_id in adapter.replaced_streaming_prompt_ids
    assert adapter._max_model_len == 4096
    assert running_queue == []
    assert waiting_queue == [request]


def test_load_poll_ar_keeps_generic_prompt_extension_for_non_tts_stage(build_adapter) -> None:
    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
    request = _req("req-generic-extension", RequestStatus.RUNNING)
    request.resumable = True
    request.prompt_token_ids = [1, 2, 3]
    request._all_token_ids = [1, 2, 3, 7]
    request._output_token_ids = [7]
    request.num_prompt_tokens = 3
    request.num_computed_tokens = 4
    request.update_block_hashes = Mock()
    adapter.get_req_chunk[request.request_id] = 1
    connector.get.return_value = (
        {
            "ids": {"prompt": [1]},
            "meta": {"next_stage_prompt_len": 2},
        },
        1,
    )

    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
    adapter.requests_with_ready_chunks.add(request.request_id)
    adapter._apply_pending_ar_prompt_updates({request.request_id: request})

    assert request.num_computed_tokens == 4
    assert request.prompt_token_ids == [1, 2, 3, 7, 0, 0]
    assert request._all_token_ids == [1, 2, 3, 7, 0, 0]
    assert request.num_prompt_tokens == 6


def test_window_condition_sequence_ignores_control_only_chunks(build_adapter) -> None:
    adapter, connector = build_adapter(
        stage_id=1,
        model_mode="ar",
        tts_attention_type="sliding_recompute",
    )
    request = _req("req-condition-seq", RequestStatus.RUNNING)

    def receive(payload):
        connector.get.return_value = (payload, 1)
        assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
        adapter.requests_with_ready_chunks.add(request.request_id)
        adapter._apply_pending_ar_prompt_updates({request.request_id: request})
        adapter.requests_with_ready_chunks.discard(request.request_id)
        return request.additional_information

    first = receive(
        {
            "native_duplex": True,
            "ids": {"prompt": [1]},
            "meta": {"next_stage_prompt_len": 10},
        }
    )
    assert first["meta"]["streaming_condition_seq"] == 0
    assert adapter._streaming_condition_seqs[request.request_id] == 0

    control = receive({"meta": {"is_segment_finished": True}})
    assert "streaming_condition_seq" not in control["meta"]
    assert adapter._streaming_condition_seqs[request.request_id] == 0

    second = receive(
        {
            "native_duplex": True,
            "ids": {"prompt": [2]},
            "meta": {"next_stage_prompt_len": 12},
        }
    )
    assert second["meta"]["streaming_condition_seq"] == 1
    assert adapter._streaming_condition_lengths[request.request_id] == 12


def test_load_poll_ar_reports_invalid_capacity_metadata(build_adapter):
    adapter, connector = build_adapter(stage_id=1, model_mode="ar", max_model_len=4096)
    request = _req("req-invalid-capacity", RequestStatus.RUNNING)
    request.resumable = True
    adapter.get_req_chunk[request.request_id] = 1
    connector.get.return_value = (
        {
            "ids": {"prompt": [1]},
            "meta": {
                "next_stage_prompt_len": 4096,
                "next_stage_generation_tokens": 26,
                "replace_streaming_prompt": True,
                "finished": False,
            },
        },
        1,
    )

    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
    request.status = RequestStatus.WAITING_FOR_CHUNK
    running_queue = [request]
    adapter.process_pending_chunks(
        DummyWaitingQueue(),
        running_queue,
        scheduler_requests={request.request_id: request},
    )
    assert adapter.collect_failed_receive_request_ids() == {
        request.request_id: "fresh streaming prompt plus generation reserve exceeds max_model_len: "
        "prompt=4096, reserve=26, limit=4096"
    }
    assert request.request_id not in adapter._finished_load_reqs


def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter):
    adapter, connector = build_adapter(stage_id=1, model_mode="generation")
    request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor")

    codes = torch.tensor([[1, 2], [3, 4]], dtype=torch.long)
    payload: OmniPayload = {
        "codes": {"audio": codes},
        "meta": {
            "left_context_size": 1,
            "finished": torch.tensor(False, dtype=torch.bool),
        },
    }
    connector.get.return_value = (payload, 16)

    adapter._poll_single_request(_dequeue_load_entry(adapter, request))

    assert request.prompt_token_ids == [0]
    assert request.num_computed_tokens == 0
    assert torch.equal(request.additional_information["codes"]["audio"], codes)
    assert request.additional_information["meta"]["left_context_size"] == 1
    assert "finished" not in request.additional_information["meta"]
    assert "req-tensor" in adapter._finished_load_reqs


def test_load_poll_generation_empty_nonterminal_chunk_keeps_polling(build_adapter):
    adapter, connector = build_adapter(stage_id=1, model_mode="generation")
    request = _req("req-empty-tensor", RequestStatus.WAITING, external_req_id="external-empty")

    empty_payload: OmniPayload = {
        "codes": {"audio": torch.empty((4, 0), dtype=torch.long)},
        "meta": {
            "left_context_size": 0,
            "finished": torch.tensor(False, dtype=torch.bool),
        },
    }
    ready_payload: OmniPayload = {
        "codes": {"audio": torch.tensor([[1, 2]], dtype=torch.long)},
        "meta": {
            "left_context_size": 0,
            "finished": torch.tensor(False, dtype=torch.bool),
        },
    }
    connector.get.side_effect = [(empty_payload, 16), (ready_payload, 16)]

    entry = _dequeue_load_entry(adapter, request)
    assert adapter._poll_single_request(entry) is False
    assert request.request_id not in adapter._finished_load_reqs
    assert request.request_id not in adapter.requests_with_ready_chunks
    assert adapter.get_req_chunk[request.request_id] == 1

    assert adapter._poll_single_request(entry) is True
    assert request.request_id in adapter._finished_load_reqs
    assert torch.equal(request.additional_information["codes"]["audio"], ready_payload["codes"]["audio"])
    assert adapter.get_req_chunk[request.request_id] == 2


def test_save_async(build_adapter):
    adapter, _ = build_adapter(stage_id=1)
    request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")

    adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": False}
    adapter.save_async(multimodal_output=None, request=request)
    adapter.custom_process_next_stage_input_func = lambda **kwargs: {}
    adapter.save_async(multimodal_output=None, request=request)

    task = adapter._pending_save_reqs.popleft()
    assert task["is_finished"] is False


def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_adapter):
    adapter, _ = build_adapter(stage_id=1)
    request = _req("req-async", RequestStatus.WAITING, external_req_id="external-async")
    request.num_computed_tokens = 10
    request.num_output_placeholders = 2

    adapter.save_async(multimodal_output=None, request=request)

    assert adapter.requests_num_chunks_sent["external-async"] == 8
    assert len(adapter._pending_save_reqs) == 1


def test_segment_boundary_starts_new_send_watermark_before_background_flush(build_adapter):
    """A queued boundary owns the end of its deduplication generation.

    The next segment can start before the save thread sends the old boundary.
    Sending that old task later must not erase the new segment's watermark.
    """
    adapter, _ = build_adapter(stage_id=1)
    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
    request.resumable = True
    request._omni_segment_generation = 1
    request.num_computed_tokens = 0
    adapter.requests_num_chunks_sent["ext-stream"] = 26

    adapter.save_async(
        multimodal_output=None,
        request=request,
        is_segment_finished=True,
        confirmed_num_computed_tokens=26,
        segment_generation=0,
    )

    assert len(adapter._pending_save_reqs) == 1
    boundary_task = adapter._pending_save_reqs.popleft()
    assert boundary_task["segment_generation"] == 0
    assert "ext-stream" not in adapter.requests_num_chunks_sent

    request.num_computed_tokens = 3
    request._omni_segment_generation = 1
    adapter.save_async(
        multimodal_output=None,
        request=request,
        is_segment_finished=False,
    )

    assert len(adapter._pending_save_reqs) == 1
    assert adapter.requests_num_chunks_sent["ext-stream"] == 3

    adapter._send_single_request(boundary_task)

    assert adapter.requests_num_chunks_sent["ext-stream"] == 3


def test_background_send_uses_enqueued_request_snapshot(build_adapter):
    """A queued segment must not observe later in-place request mutations."""
    adapter, _ = build_adapter(stage_id=1)
    request = Request(
        request_id="req-stream",
        prompt_token_ids=[1, 2],
        sampling_params=SamplingParams(max_tokens=8),
        pooling_params=None,
        resumable=True,
    )
    request.external_req_id = "ext-stream"
    ref_audio = torch.tensor([0.1, -0.1])
    request.additional_information = {"codes": {"ref": ref_audio}, "meta": {"segment": "old"}}
    request.append_output_token_ids([7])
    seen_requests = []

    def recording_processor(**kwargs):
        queued = kwargs["request"]
        seen_requests.append(
            (
                queued.additional_information["meta"]["segment"],
                list(queued.prompt_token_ids),
                list(queued.output_token_ids),
                list(queued.all_token_ids),
                queued.additional_information["codes"]["ref"] is ref_audio,
            )
        )
        return OmniPayloadStruct()

    adapter.custom_process_next_stage_input_func = recording_processor
    adapter.save_async(
        multimodal_output=None,
        request=request,
        is_segment_finished=True,
    )
    request.additional_information["meta"]["segment"] = "next"
    request.prompt_token_ids.append(3)
    request.append_output_token_ids([8])
    request._omni_segment_generation = 1
    adapter.save_async(
        multimodal_output=None,
        request=request,
        is_segment_finished=True,
    )
    request.additional_information["meta"]["segment"] = "later"
    request.prompt_token_ids.append(4)
    request.append_output_token_ids([9])

    assert len(adapter._pending_save_reqs) == 2
    first_task = adapter._pending_save_reqs.popleft()
    second_task = adapter._pending_save_reqs.popleft()
    assert first_task["request"].additional_information is not request.additional_information
    assert first_task["request"].additional_information["meta"] is not request.additional_information["meta"]

    adapter._send_single_request(first_task)
    adapter._send_single_request(second_task)

    assert seen_requests == [
        ("old", [1, 2], [7], [1, 2, 7], True),
        ("next", [1, 2, 3], [7, 8], [1, 2, 7, 8], True),
    ]


def test_save_without_custom_processor_does_not_snapshot_request(build_adapter, mocker):
    adapter, _ = build_adapter(stage_id=1)
    request = _req("req-direct", RequestStatus.RUNNING, external_req_id="ext-direct")
    snapshot = mocker.patch.object(
        adapter,
        "_snapshot_processor_request",
        side_effect=AssertionError("unexpected snapshot"),
    )

    adapter.save_async(multimodal_output=None, request=request)

    snapshot.assert_not_called()
    assert adapter._pending_save_reqs.popleft()["request"] is request


def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch):
    """A terminal stop is not a segment boundary (#5383), but the producer-side
    processor must still receive the flush signal on the terminal chunk.
    Passing only ``is_segment_finished`` starved processors of their final
    accumulated payload once terminal stops stopped setting it (#5413: the
    downstream stage got ``meta.finished`` with the tail data missing).
    """
    adapter, connector = build_adapter(stage_id=0)
    request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal")

    seen_flush_flags = []

    def recording_processor(**kwargs):
        seen_flush_flags.append(kwargs["is_finished"])
        return OmniPayloadStruct(
            codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.long)),
        )

    adapter.custom_process_next_stage_input_func = recording_processor
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)

    adapter._send_single_request(
        {"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": False}
    )

    assert seen_flush_flags == [True]
    sent_payload = connector.put.call_args.kwargs["data"]
    assert bool(sent_payload.meta.finished.item()) is True
    assert bool(sent_payload.meta.is_segment_finished.item()) is False


def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch):
    """Producer may return a struct with ``meta=None`` (e.g. payload that
    carries only ``embed`` or ``codes``). The sender's ``meta is not None``
    guard handles this without AttributeError; ``finished_flag`` is None and
    the cleanup path is not triggered.
    """
    adapter, _ = build_adapter(stage_id=1)
    request = _req("req-no-meta", RequestStatus.WAITING, external_req_id="ext-no-meta")

    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
        codes=CodesStruct(audio=torch.tensor([1, 2], dtype=torch.long)),
    )
    cleanup_calls = []
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))

    adapter._send_single_request(
        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
    )

    assert cleanup_calls == []  # no terminal cleanup; meta.finished is false


def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatch):
    """Pin the contract: an explicitly empty ``OmniPayloadStruct()`` passes
    the ``payload_data is None`` check and gets sent. To skip a chunk, the
    producer must return ``None``, not an empty struct. (Filtering empty
    structs at the adapter would require introspecting all struct fields on
    every send and was rejected for cost vs. value.)
    """
    adapter, connector = build_adapter(stage_id=1)
    request = _req("req-empty", RequestStatus.WAITING, external_req_id="ext-empty")

    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)

    adapter._send_single_request(
        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
    )

    assert connector.put.called
    sent_payload = connector.put.call_args.kwargs["data"]
    assert isinstance(sent_payload, OmniPayloadStruct)
    assert sent_payload.meta.finished.item() is False
    assert sent_payload.meta.is_segment_finished.item() is False


def test_send_single_request_struct_preserves_segment_finished(build_adapter, monkeypatch):
    adapter, connector = build_adapter(stage_id=1)
    request = _req("req-segment", RequestStatus.WAITING, external_req_id="ext-segment")

    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)

    adapter._send_single_request(
        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
    )

    sent_payload = connector.put.call_args.kwargs["data"]
    assert sent_payload.meta.finished.item() is False
    assert sent_payload.meta.is_segment_finished.item() is True


def test_send_single_request_respects_processor_receiver_boundary(build_adapter, monkeypatch):
    adapter, connector = build_adapter(stage_id=1)
    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")

    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
        meta=MetaStruct(is_segment_finished=torch.tensor(False, dtype=torch.bool))
    )
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)

    adapter._send_single_request(
        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
    )

    sent_payload = connector.put.call_args.kwargs["data"]
    assert sent_payload.meta.is_segment_finished.item() is False


def test_send_single_request_personaplex_pending_frame_is_not_segment_boundary(
    build_adapter,
):
    from vllm_omni.model_executor.stage_input_processors.personaplex import (
        talker2code2wav_async_chunk,
    )

    adapter, connector = build_adapter(
        stage_id=0,
        connector_extra={
            "initial_codec_chunk_frames": 1,
            "codec_chunk_frames": 5,
        },
    )
    request = _req(
        "req-personaplex",
        RequestStatus.WAITING,
        external_req_id="ext-personaplex",
    )
    request.resumable = True
    request.additional_information = {
        "codes": {
            "audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
        }
    }
    adapter.custom_process_next_stage_input_func = talker2code2wav_async_chunk

    adapter._send_single_request(
        {
            "multimodal_output": None,
            "request": request,
            "is_finished": False,
            "is_segment_finished": True,
        }
    )

    sent_payload = connector.put.call_args.kwargs["data"]
    assert sent_payload.codes is None
    assert sent_payload.meta.finished.item() is False
    assert sent_payload.meta.is_segment_finished.item() is False


def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter):
    from vllm_omni.model_executor.stage_input_processors.personaplex import (
        talker2code2wav_async_chunk,
    )

    adapter, _ = build_adapter(
        stage_id=0,
        connector_extra={
            "initial_codec_chunk_frames": 1,
            "codec_chunk_frames": 5,
        },
    )
    request = _req(
        "req-personaplex-first",
        RequestStatus.WAITING,
        external_req_id="ext-personaplex-reused",
    )
    request.resumable = True
    request.additional_information = {
        "codes": {
            "audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
        }
    }

    first = talker2code2wav_async_chunk(
        adapter,
        multimodal_output=None,
        request=request,
        is_finished=True,
    )
    assert first is not None
    assert first.codes is None

    adapter.cleanup_sender(request.external_req_id)

    replacement = _req(
        "req-personaplex-replacement",
        RequestStatus.WAITING,
        external_req_id=request.external_req_id,
    )
    replacement.resumable = True
    replacement.additional_information = {
        "codes": {
            "audio": torch.arange(8, 16, dtype=torch.long).reshape(1, 8),
        }
    }
    second = talker2code2wav_async_chunk(
        adapter,
        multimodal_output=None,
        request=replacement,
        is_finished=True,
    )

    assert second is not None
    assert second.codes is None


def test_save_async_skips_stale_resumable_chunk_within_segment(build_adapter):
    adapter, _ = build_adapter(stage_id=1)
    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
    request.resumable = True
    request.num_computed_tokens = 0
    adapter.requests_num_chunks_sent["ext-stream"] = 111

    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)

    assert len(adapter._pending_save_reqs) == 0
    assert adapter.requests_num_chunks_sent["ext-stream"] == 111

    adapter.requests_num_chunks_sent.pop("ext-stream")
    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)

    assert len(adapter._pending_save_reqs) == 1
    assert adapter.requests_num_chunks_sent["ext-stream"] == 0


def test_save_async_drops_late_previous_segment_after_boundary_reset(build_adapter):
    adapter, _ = build_adapter(stage_id=1)
    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
    request.resumable = True
    request._omni_segment_generation = 0
    adapter.requests_num_chunks_sent["ext-stream"] = 111

    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)

    late = _req("req-stream-late", RequestStatus.WAITING, external_req_id="ext-stream")
    late.resumable = True
    late.num_computed_tokens = 111
    late._omni_segment_generation = 0
    adapter.save_async(multimodal_output=None, request=late, is_segment_finished=False)

    next_segment = _req("req-stream-next", RequestStatus.WAITING, external_req_id="ext-stream")
    next_segment.resumable = True
    next_segment.num_computed_tokens = 3
    next_segment._omni_segment_generation = 1
    adapter.save_async(multimodal_output=None, request=next_segment, is_segment_finished=False)

    assert len(adapter._pending_save_reqs) == 2
    assert adapter.requests_num_chunks_sent["ext-stream"] == 3
    assert adapter._pending_save_reqs[-1]["request"] is next_segment


def test_send_single_request_preserves_queued_preboundary_chunk(build_adapter):
    adapter, connector = build_adapter(stage_id=1)
    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(meta=MetaStruct())
    request = _req("req-queued-old", RequestStatus.WAITING, external_req_id="ext-queued-old")
    request.resumable = True
    request._omni_segment_generation = 0

    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)

    old_task = adapter._pending_save_reqs.popleft()
    adapter._send_single_request(old_task)
    connector.put.assert_called_once()

    boundary_task = adapter._pending_save_reqs.popleft()
    adapter._send_single_request(boundary_task)
    assert connector.put.call_count == 2


def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
    adapter, _ = build_adapter(stage_id=1)
    request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")

    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
        meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool))
    )
    cleanup_calls = []
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))

    adapter._send_single_request(
        {"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
    )

    assert len(cleanup_calls) == 1
    args, _ = cleanup_calls[0]
    assert args[0] == "req-finished"
    assert args[1] == "ext-finished"


def test_old_turn_terminal_send_preserves_new_turn_state_and_sender(build_adapter):
    adapter, connector = build_adapter(
        stage_id=1,
        model_mode="ar",
        max_model_len=4096,
        tts_attention_type="sliding_recompute",
    )
    request = _req("req-live", RequestStatus.RUNNING, external_req_id="ext-live")
    request.resumable = True
    request._all_token_ids = []
    request._output_token_ids = []
    request.update_block_hashes = Mock()
    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
        meta=MetaStruct(
            finished=torch.tensor(True, dtype=torch.bool),
            codec_streaming=True,
        )
    )

    put_started = threading.Event()
    release_put = threading.Event()
    put_keys = []

    def blocking_first_put(*_args, **kwargs):
        put_keys.append(kwargs["put_key"])
        if len(put_keys) == 1:
            put_started.set()
            assert release_put.wait(timeout=5)
        return True, 1, {}

    connector.put.side_effect = blocking_first_put
    adapter.save_async(multimodal_output=None, request=request)
    first_task = adapter._pending_save_reqs.popleft()
    first_sender_token = first_task["sender_token"]
    send_thread = threading.Thread(target=lambda: adapter._send_single_request(first_task))
    send_thread.start()
    assert put_started.wait(timeout=5)

    entry = _dequeue_load_entry(adapter, request)
    connector.get.return_value = (
        {
            "ids": {"prompt": [1]},
            "meta": {
                "replace_streaming_prompt": True,
                "next_stage_prompt_len": 10,
                "next_stage_generation_tokens": 26,
            },
        },
        1,
    )
    assert adapter._poll_single_request(entry) is True
    adapter.requests_with_ready_chunks.add(request.request_id)
    adapter._apply_pending_ar_prompt_updates({request.request_id: request})
    assert adapter._streaming_condition_seqs[request.request_id] == 0

    request.num_computed_tokens = 1
    adapter.save_async(multimodal_output=None, request=request)
    second_task = adapter._pending_save_reqs.popleft()
    assert second_task["sender_token"] is first_sender_token

    release_put.set()
    send_thread.join(timeout=5)
    assert not send_thread.is_alive()
    assert adapter._streaming_condition_seqs[request.request_id] == 0
    assert adapter._sender_tokens[request.external_req_id] is first_sender_token
    assert first_sender_token.cancelled is False

    adapter._send_single_request(second_task)
    assert put_keys == ["ext-live_1_0", "ext-live_1_1"]


def test_load_poll_non_ar_merges_into_existing_additional_information(build_adapter):
    adapter, connector = build_adapter(stage_id=2, model_mode="diffusion")
    request = _req("req-non-ar", RequestStatus.WAITING, external_req_id="ext-non-ar")
    request.additional_information = {
        "hidden_states": {"output": torch.tensor([[1.0]])},
        "ids": {"prompt": [11, 12]},
        "meta": {"finished": torch.tensor(False, dtype=torch.bool), "step": 1},
    }
    request.num_computed_tokens = 9

    payload: OmniPayload = {
        "hidden_states": {"output": torch.tensor([[2.0]])},
        "ids": {"all": [21, 22]},
        "codes": {"audio": torch.tensor([7, 8], dtype=torch.long)},
        "meta": {"finished": torch.tensor(True, dtype=torch.bool), "phase": "decode"},
        "kv_metadata": {"foo": "bar"},
    }
    connector.get.return_value = (payload, 8)

    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True

    assert request.prompt_token_ids == [7, 8]
    assert request.num_computed_tokens == 0
    assert torch.equal(
        request.additional_information["hidden_states"]["output"],
        torch.tensor([[2.0]]),
    )
    assert request.additional_information["ids"]["prompt"] == [11, 12]
    assert request.additional_information["ids"]["all"] == [21, 22]
    # non-ar merge path intentionally doesn't overwrite meta.finished.
    assert request.additional_information["meta"]["finished"].item() is False
    assert request.additional_information["meta"]["phase"] == "decode"
    assert request.additional_information["kv_metadata"] == {"foo": "bar"}
    assert "req-non-ar" in adapter._finished_load_reqs
    assert "req-non-ar" in adapter.upstream_exhausted_requests


def test_load_poll_generation_segment_marker_replaces_previous_chunk(build_adapter):
    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
    request = _req("req-marker", RequestStatus.WAITING, external_req_id="external-marker")
    request.additional_information = {
        "codes": {"audio": torch.tensor([1, 2])},
        "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
    }
    connector.get.return_value = (
        {
            "codes": {
                "audio": torch.tensor([7, 8], dtype=torch.long),
                "ref": torch.tensor([0.1, -0.1]),
            },
            "meta": {
                "finished": torch.tensor(False, dtype=torch.bool),
                "is_segment_finished": torch.tensor(True, dtype=torch.bool),
                "request_finished": torch.tensor(False, dtype=torch.bool),
                "replace_runtime_additional_information": True,
            },
        },
        1,
    )

    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True

    assert request.prompt_token_ids == [7, 8]
    assert "audio" not in request.additional_information["codes"]
    torch.testing.assert_close(
        request.additional_information["codes"]["ref"],
        torch.tensor([0.1, -0.1]),
    )
    assert not {"cache_epoch", "chunk_seq", "last_chunk"}.intersection(request.additional_information["meta"])
    assert request.request_id in adapter.segment_finished_requests


@pytest.mark.parametrize("cache_epoch", [0, 1])
def test_generation_first_chunk_replaces_prewarm_runner_payload(build_adapter, cache_epoch):
    from vllm_omni.core.sched.output import OmniNewRequestData

    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
    request = _req("req-first-chunk", RequestStatus.WAITING)
    request.model_intermediate_buffer = {"duplex": {"data_plane": True}, "global_request_id": "session"}
    connector.get.return_value = (
        {
            "codes": {"audio": torch.tensor([7, 8, 9, 10])},
            "meta": {
                "cache_epoch": cache_epoch,
                "chunk_seq": 0,
                "code_flat_numel": 4,
                "replace_runtime_additional_information": True,
            },
        },
        1,
    )

    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
    scheduled = OmniNewRequestData.from_request(request, block_ids=([0],))

    # New/resumed requests prioritize model_intermediate_buffer in the runner;
    # it must carry the same codec snapshot as cached-request updates.
    assert scheduled.model_intermediate_buffer == scheduled.additional_information
    assert scheduled.model_intermediate_buffer["meta"]["cache_epoch"] == cache_epoch
    assert scheduled.model_intermediate_buffer["meta"]["chunk_seq"] == 0
    assert scheduled.prompt_token_ids == [7, 8, 9, 10]


def test_load_poll_generation_empty_replacement_snapshot_is_ready(build_adapter):
    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
    request = _req("req-empty-marker", RequestStatus.WAITING, external_req_id="external-empty-marker")
    request.additional_information = {
        "codes": {"audio": torch.tensor([1, 2])},
        "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
    }
    connector.get.return_value = (
        {
            "meta": {
                "is_segment_finished": torch.tensor(True, dtype=torch.bool),
                "replace_runtime_additional_information": True,
            }
        },
        1,
    )

    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True

    assert request.prompt_token_ids == [0]
    assert "codes" not in request.additional_information
    assert request.additional_information["meta"]["replace_runtime_additional_information"] is True
    assert request.request_id in adapter.segment_finished_requests
    assert request.request_id in adapter._finished_load_reqs


def test_load_poll_generation_without_snapshot_marker_keeps_incremental_state(build_adapter):
    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
    request = _req("req-incremental", RequestStatus.WAITING, external_req_id="external-incremental")
    request.additional_information = {
        "codes": {"audio": torch.tensor([1, 2])},
        "meta": {"cache_epoch": 3, "chunk_seq": 2},
    }
    connector.get.return_value = (
        {
            "meta": {
                "finished": torch.tensor(False, dtype=torch.bool),
                "phase": "decode",
            }
        },
        1,
    )

    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is False

    assert torch.equal(request.additional_information["codes"]["audio"], torch.tensor([1, 2]))
    assert request.additional_information["meta"]["cache_epoch"] == 3
    assert request.additional_information["meta"]["chunk_seq"] == 2
    assert request.additional_information["meta"]["phase"] == "decode"


def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
    request = _req("req-merged", RequestStatus.WAITING, external_req_id="ext-merged")
    request.additional_information = {
        "hidden_states": {"output": torch.tensor([[1.0]])},
        "ids": {"prompt": [11, 12]},
        "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
    }

    adapter.request_ids_mapping["req-merged"] = "ext-merged"
    payload: OmniPayload = {
        "hidden_states": {"output": torch.tensor([[2.0]])},
        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
    }
    connector.get.return_value = (payload, 8)

    adapter._poll_single_request(_dequeue_load_entry(adapter, request))

    # The recv thread records the latest payload; scheduler-side queue
    # restoration publishes it onto the Request.
    assert torch.equal(request.additional_information["hidden_states"]["output"], torch.tensor([[1.0]]))
    request.status = RequestStatus.WAITING_FOR_CHUNK
    adapter.process_pending_chunks(
        DummyWaitingQueue(),
        [request],
        scheduler_requests={request.request_id: request},
    )
    assert request.additional_information == payload
    assert request.additional_information["meta"]["finished"].item() is True


def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter):
    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
    request = _req("req-later-chunk", RequestStatus.WAITING, external_req_id="ext-later-chunk")
    request.prefill_stats = PrefillStats()
    adapter.request_ids_mapping[request.request_id] = request.external_req_id

    connector.get.return_value = (
        {
            "codes": {"audio": torch.tensor([7, 8], dtype=torch.long)},
            "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
        },
        8,
    )
    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
    OmniGenerationScheduler._record_prefill_stats(request)
    first_chunk_stats = request.prefill_stats
    request.prefill_stats = None

    connector.get.return_value = (
        {
            "codes": {"audio": torch.tensor([9, 10, 11], dtype=torch.long)},
            "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
        },
        8,
    )
    assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
    assert isinstance(request.prefill_stats, PrefillStats)
    OmniGenerationScheduler._record_prefill_stats(request)
    second_chunk_stats = request.prefill_stats

    assert first_chunk_stats is not None
    assert first_chunk_stats.num_prompt_tokens == 2
    assert second_chunk_stats is not None
    assert second_chunk_stats.num_prompt_tokens == 3

    prompt_token_stats = PromptTokenStats()
    prompt_token_stats.update_from_output(first_chunk_stats)
    prompt_token_stats.update_from_output(second_chunk_stats)
    assert prompt_token_stats.total == 5
    assert prompt_token_stats.computed == 5


def test_sender_only_adapter_does_not_park_or_clear_requests(build_adapter):
    adapter, _ = build_adapter(stage_id=1, connector_extra={"role": "sender"})
    request = _req("req-1", RequestStatus.WAITING)
    request.additional_information = {"tts_token_ids": torch.tensor([1])}
    waiting_queue = DummyWaitingQueue([request])
    running_queue: list = []

    adapter.load_async(request)
    adapter.process_pending_chunks(
        waiting_queue,
        running_queue,
        scheduler_requests={request.request_id: request},
    )

    assert waiting_queue == [request]
    assert request.status == RequestStatus.WAITING
    assert request.additional_information["tts_token_ids"].item() == 1
    assert adapter._pending_load_reqs == deque()


def test_process_and_restore_queues(build_adapter):
    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
    waiting_req = _req("w1", RequestStatus.WAITING)
    running_req = _req("r1", RequestStatus.RUNNING)
    waiting_queue = DummyWaitingQueue([waiting_req])
    running_queue = [running_req]
    scheduler_requests = {waiting_req.request_id: waiting_req, running_req.request_id: running_req}

    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
    assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK
    assert running_req.status == RequestStatus.WAITING_FOR_CHUNK
    assert waiting_queue == []
    assert running_queue == []

    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
    assert waiting_queue == [waiting_req]
    assert running_queue == [running_req]
    assert adapter.waiting_for_chunk_waiting_requests == deque()
    assert adapter.waiting_for_chunk_running_requests == deque()


def test_fifo_promotion(build_adapter):
    adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
    reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
    waiting_queue = DummyWaitingQueue(reqs)
    running_queue: list = []

    adapter.process_pending_chunks(waiting_queue, running_queue)

    assert list(adapter._active_streams) == ["req-1", "req-2"]
    assert waiting_queue == []
    assert reqs[0].status == RequestStatus.WAITING_FOR_CHUNK
    assert reqs[1].status == RequestStatus.WAITING_FOR_CHUNK
    assert reqs[2].status == RequestStatus.WAITING

    # Ordinary completion calls cleanup_receiver(), not finish_requests()
    # (the abort path) -- see test_omni_ar_scheduler_free_request_cleanup.py.
    # Eviction is deferred to postprocess_scheduler_output in the runtime path
    # (commit c4d95fd9 -- otherwise the terminal chunk deadlocks at c=8 K=2).
    # Simulate the restore step here so promotion can pick up the freed slot.
    adapter.cleanup_receiver("req-1")
    adapter.restore_queues(waiting_queue, running_queue)
    adapter.process_pending_chunks(waiting_queue, running_queue)

    assert list(adapter._active_streams) == ["req-2", "req-3"]
    # Promotion + chunk processing happen in the same call, so req-3 is
    # already WAITING_FOR_CHUNK by the time we check.
    assert reqs[2].status == RequestStatus.WAITING_FOR_CHUNK


def test_non_active_waiting_request_is_held_off_scheduler(build_adapter):
    adapter, _ = build_adapter(stage_id=2, max_num_seqs=2, active_stream_window=1)
    active = _req("req-active", RequestStatus.WAITING)
    non_active = _req("req-non-active", RequestStatus.WAITING)
    waiting_queue = DummyWaitingQueue([active, non_active])
    running_queue: list = []

    adapter.process_pending_chunks(waiting_queue, running_queue)

    assert waiting_queue == []
    assert active.status == RequestStatus.WAITING_FOR_CHUNK
    assert non_active.status == RequestStatus.WAITING
    assert [entry.request for entry in adapter._pending_load_reqs] == [active]
    assert list(adapter.waiting_for_chunk_waiting_requests) == [active, non_active]


def test_legacy_k0(build_adapter):
    adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=0)
    waiting_req = _req("waiting", RequestStatus.WAITING)
    running_req_1 = _req("running-1", RequestStatus.RUNNING)
    running_req_2 = _req("running-2", RequestStatus.RUNNING)
    waiting_queue = DummyWaitingQueue([waiting_req])
    running_queue = [running_req_1, running_req_2]

    adapter.requests_with_ready_chunks.update({"running-1", "running-2"})
    adapter.process_pending_chunks(waiting_queue, running_queue)

    assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK
    assert adapter.waiting_for_chunk_waiting_requests == deque([waiting_req])
    assert running_queue == [running_req_1]
    assert waiting_queue == [running_req_2]
    assert running_req_2.status == RequestStatus.PREEMPTED
    assert adapter._active_streams == {}


def test_finished_releases_slot(build_adapter):
    adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=1, active_stream_window=1)
    req_1 = _req("req-1", RequestStatus.WAITING)
    req_2 = _req("req-2", RequestStatus.WAITING)
    waiting_queue = DummyWaitingQueue([req_1, req_2])
    running_queue: list = []

    adapter.process_pending_chunks(waiting_queue, running_queue)
    assert list(adapter._active_streams) == ["req-1"]
    assert waiting_queue == []

    # req-1 finishes ordinarily (see test_fifo_promotion for why this is
    # cleanup_receiver(), not finish_requests()). Eviction is deferred to
    # postprocess_scheduler_output in the runtime path (commit c4d95fd9);
    # simulate the restore step so promotion can pick up the freed slot.
    adapter.cleanup_receiver("req-1")
    adapter.restore_queues(waiting_queue, running_queue)
    adapter.process_pending_chunks(waiting_queue, running_queue)

    assert list(adapter._active_streams) == ["req-2"]
    # Promotion + chunk processing happen in the same call, so req-2 is
    # already WAITING_FOR_CHUNK by the time we check.
    assert req_2.status == RequestStatus.WAITING_FOR_CHUNK


def test_postprocess_scheduler_output(build_adapter):
    adapter, _ = build_adapter()
    adapter.requests_with_ready_chunks = {"new-ready", "cached-ready", "leftover"}

    scheduler_output = SimpleNamespace(
        scheduled_new_reqs=[SimpleNamespace(req_id="new-ready")],
        scheduled_cached_reqs=SimpleNamespace(req_ids=["cached-ready", "missing"]),
    )
    requests = {"cached-ready": SimpleNamespace(additional_information={"k": "v"})}

    adapter.postprocess_scheduler_output(scheduler_output, requests)

    cached_info = scheduler_output.scheduled_cached_reqs.additional_information
    assert cached_info["cached-ready"] == {"k": "v"}
    assert cached_info["missing"] is None
    assert adapter.requests_with_ready_chunks == {"leftover"}


@pytest.mark.parametrize("model_mode", ["ar", "generation"])
def test_active_stream_window_stalls_lone_running_request_after_upstream_finishes(build_adapter, model_mode):
    """Regression test for vllm-project/vllm-omni#5349.

    upstream_exhausted_requests means the upstream sent its terminal chunk,
    not that this stage's own generation is done. A downstream stage can
    still have a long decode ahead after its upstream finishes. Evicting on
    that signal and then permanently denying re-admission wedges the
    request out of running_queue forever, with no error.

    requests_with_ready_chunks is set up front so the request stays in
    running_queue this tick instead of legitimately parking in
    waiting_for_chunk_running_requests (a separate mechanism this test must
    not conflate with the bug).

    Parametrized over both worker types: the real issue hit an "ar" stage
    (Qwen3-Omni's talker); must not regress "generation" (Qwen3-TTS's
    Code2Wav), which this feature was designed for.
    """
    adapter, _ = build_adapter(stage_id=1, model_mode=model_mode, max_num_seqs=2, active_stream_window=2)
    running_req = _req("req-1", RequestStatus.RUNNING)
    running_queue = [running_req]
    waiting_queue = DummyWaitingQueue([])

    adapter._active_streams["req-1"] = running_req
    adapter.requests_with_ready_chunks.add("req-1")

    # Upstream terminal chunk arrives; this stage hasn't finished.
    adapter.upstream_exhausted_requests.add("req-1")
    scheduler_output = SimpleNamespace(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=SimpleNamespace(req_ids=["req-1"]),
    )
    adapter.postprocess_scheduler_output(scheduler_output)
    assert not running_req.is_finished()  # sanity: not done

    # No competition for the slot; req-1 must stay schedulable.
    adapter.process_pending_chunks(waiting_queue, running_queue)

    assert running_queue == [running_req], "req-1 must still be schedulable -- losing it here is #5349"
    assert not adapter.waiting_for_chunk_running_requests
    assert running_req not in adapter._held_non_active
    assert list(adapter._active_streams) == ["req-1"]


def test_cleanup_receiver_releases_multiple_slots_in_sequence(build_adapter):
    """Companion to test_fifo_promotion: after EACH of the K active streams
    finishes (not just one), promotion of a new K-sized batch must still
    work -- the window must not stay exhausted by stale entries. See
    test_omni_ar_scheduler_free_request_cleanup.py for the scheduler-level
    proof that ordinary completion calls cleanup_receiver().
    """
    adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
    reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
    waiting_queue = DummyWaitingQueue(reqs)
    running_queue: list = []

    adapter.process_pending_chunks(waiting_queue, running_queue)
    assert list(adapter._active_streams) == ["req-1", "req-2"]

    adapter.cleanup_receiver("req-1")
    adapter.cleanup_receiver("req-2")
    # Eviction is deferred to postprocess_scheduler_output in the runtime
    # path; simulate the restore step so promotion can pick up the freed
    # slots (see test_fifo_promotion).
    adapter.restore_queues(waiting_queue, running_queue)
    adapter.process_pending_chunks(waiting_queue, running_queue)

    assert list(adapter._active_streams) == ["req-3", "req-4"], (
        "both freed slots must be reusable after ordinary completion"
    )


# ---------------------------------------------------------------
# Cleanup tests
# ---------------------------------------------------------------


def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"):
    """Fill every per-request structure so cleanup can be verified."""
    adapter.upstream_exhausted_requests.add(req_id)
    adapter._active_streams[req_id] = SimpleNamespace(request_id=req_id)
    adapter.get_req_chunk[req_id] = 3
    adapter.requests_with_ready_chunks.add(req_id)
    adapter.request_ids_mapping[req_id] = ext_id
    entry = SimpleNamespace(request_id=req_id)
    adapter._pending_load_reqs.append(entry)
    adapter._registered_load_entries[req_id] = entry
    adapter._finished_load_reqs.add(req_id)
    adapter._pending_ar_prompt_updates[req_id] = (
        SimpleNamespace(request_id=req_id),
        {},
        False,
        False,
        False,
    )
    adapter._streaming_condition_lengths[req_id] = 10
    adapter._streaming_condition_seqs[req_id] = 2

    adapter.put_req_chunk[ext_id] = 5
    adapter.request_payload[ext_id] = {"hidden": [1, 2]}
    adapter.code_prompt_token_ids[ext_id] = [[10, 20]]


def test_cleanup_clears_all_state(build_adapter):
    """After cleanup, no per-request key should remain in any dict/set."""
    adapter, _ = build_adapter(stage_id=1)
    req_id, ext_id = "req-1", "ext-1"
    _populate_adapter_state(adapter, req_id, ext_id)

    adapter.cleanup(req_id, ext_id)

    assert req_id not in adapter.upstream_exhausted_requests
    assert req_id not in adapter._active_streams
    assert req_id not in adapter.get_req_chunk
    assert req_id not in adapter.requests_with_ready_chunks
    assert req_id not in adapter.request_ids_mapping
    assert req_id not in adapter._registered_load_entries
    assert req_id not in adapter._finished_load_reqs
    assert req_id not in adapter._pending_ar_prompt_updates
    assert req_id not in adapter._streaming_condition_lengths
    assert req_id not in adapter._streaming_condition_seqs
    assert req_id not in adapter._registered_load_entries

    assert ext_id not in adapter.put_req_chunk
    assert ext_id not in adapter.request_payload
    assert ext_id not in adapter.code_prompt_token_ids


def test_cleanup_and_reregister_discards_inflight_old_segment_chunk(build_adapter):
    """A late old-segment chunk cannot overwrite a resumed registration."""
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    request = _req("req-late", RequestStatus.RUNNING, external_req_id="ext-late")
    request.resumable = True
    entry = _dequeue_load_entry(adapter, request)

    get_started = threading.Event()
    release_get = threading.Event()

    def blocking_get(*_args):
        get_started.set()
        assert release_get.wait(timeout=5)
        return (
            {
                "ids": {"prompt": [1]},
                "meta": {
                    "next_stage_prompt_len": 10,
                },
            },
            1,
        )

    connector.get.side_effect = blocking_get
    poll_results = []
    poll_thread = threading.Thread(target=lambda: poll_results.append(adapter._poll_single_request(entry)))
    poll_thread.start()
    assert get_started.wait(timeout=5)

    adapter.cleanup_receiver(request.request_id)
    fresh_entry = _dequeue_load_entry(adapter, request)
    release_get.set()
    poll_thread.join(timeout=5)

    assert not poll_thread.is_alive()
    assert poll_results == [True]
    assert request.additional_information is None
    assert request.request_id not in adapter._pending_ar_prompt_updates
    assert request.request_id not in adapter._streaming_condition_seqs
    assert adapter.get_req_chunk[request.request_id] == 0
    assert adapter._registered_load_entries[request.request_id] is fresh_entry

    connector.get.side_effect = None
    connector.get.return_value = (
        {
            "ids": {"prompt": [2]},
            "meta": {"finished": True},
        },
        1,
    )
    assert adapter._poll_single_request(fresh_entry) is True
    assert adapter.get_req_chunk[request.request_id] == 1
    assert request.request_id not in adapter._registered_load_entries


def test_cleanup_infers_external_id(build_adapter):
    """When external_req_id is None, cleanup should look it up from the mapping."""
    adapter, _ = build_adapter(stage_id=1)
    req_id, ext_id = "req-2", "ext-2"
    _populate_adapter_state(adapter, req_id, ext_id)

    adapter.cleanup(req_id)

    assert ext_id not in adapter.put_req_chunk
    assert ext_id not in adapter.request_payload


def test_cleanup_idempotent(build_adapter):
    """Calling cleanup multiple times for the same (or nonexistent) request must not raise."""
    adapter, _ = build_adapter(stage_id=1)

    try:
        adapter.cleanup("nonexistent")
        adapter.cleanup("nonexistent")
    except Exception as e:
        pytest.fail(f"cleanup should be idempotent: {e}")

    req_id, ext_id = "req-3", "ext-3"
    _populate_adapter_state(adapter, req_id, ext_id)
    adapter.cleanup(req_id, ext_id)

    try:
        adapter.cleanup(req_id, ext_id)
    except Exception as e:
        pytest.fail(f"second cleanup should be idempotent: {e}")


def test_cleanup_request_id_reuse_not_polluted(build_adapter):
    """After cleanup, reusing the same request_id must not be treated as finished."""
    adapter, _ = build_adapter(stage_id=1)
    req_id, ext_id = "req-reuse", "ext-reuse"
    _populate_adapter_state(adapter, req_id, ext_id)

    adapter.cleanup(req_id, ext_id)

    assert req_id not in adapter.upstream_exhausted_requests
    assert req_id not in adapter.get_req_chunk


def test_cleanup_preserves_pending_save(build_adapter):
    """Cleanup must NOT remove _pending_save_reqs to avoid losing unsent chunks."""
    adapter, _ = build_adapter(stage_id=1)
    req_id, ext_id = "req-4", "ext-4"
    _populate_adapter_state(adapter, req_id, ext_id)

    pending_task = {"put_key": f"{ext_id}_1_0", "data": {"x": 1}}
    adapter._pending_save_reqs.append(pending_task)

    adapter.cleanup(req_id, ext_id)

    assert len(adapter._pending_save_reqs) == 1


def test_abort_invalidates_queued_sender_task_before_external_id_reuse(build_adapter):
    adapter, connector = build_adapter(stage_id=1)
    request = _req("req-old", RequestStatus.WAITING, external_req_id="ext-reused")
    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
        codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
    )

    adapter.save_async(multimodal_output=None, request=request)
    stale_task = adapter._pending_save_reqs.popleft()
    adapter.cleanup(request.request_id, request.external_req_id)
    adapter._send_single_request(stale_task)

    connector.put.assert_not_called()
    assert request.external_req_id not in adapter.put_req_chunk

    replacement = _req("req-new", RequestStatus.WAITING, external_req_id=request.external_req_id)
    adapter.save_async(multimodal_output=None, request=replacement)
    current_task = adapter._pending_save_reqs.popleft()
    adapter._send_single_request(current_task)

    connector.put.assert_called_once()
    assert connector.put.call_args.kwargs["put_key"] == "ext-reused_1_0"


def test_finish_requests_does_not_wait_for_inflight_send(build_adapter):
    adapter, connector = build_adapter(stage_id=1)
    first = _req("req-first", RequestStatus.WAITING, external_req_id="ext-first")
    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
    put_started = threading.Event()
    release_put = threading.Event()

    def blocking_put(**kwargs):
        put_started.set()
        release_put.wait(timeout=2)
        return True, 1, {}

    connector.put.side_effect = blocking_put
    adapter.save_async(multimodal_output=None, request=first)
    task = adapter._pending_save_reqs.popleft()
    sender = threading.Thread(target=adapter._send_single_request, args=(task,))
    sender.start()
    assert put_started.wait(timeout=1)
    cleaner = threading.Thread(
        target=adapter.finish_requests,
        args=([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first}),
    )
    cleaner.start()
    try:
        cleaner.join(timeout=0.5)
        assert not cleaner.is_alive()
        assert task["sender_token"].cancelled
        assert first.external_req_id in adapter._sender_tokens
    finally:
        release_put.set()
        sender.join(timeout=1)
        cleaner.join(timeout=1)
    assert not sender.is_alive()
    assert first.external_req_id not in adapter._sender_tokens
    assert first.external_req_id not in adapter.put_req_chunk


def test_cleanup_only_affects_target_request(build_adapter):
    """Cleanup for one request must not affect another request's state."""
    adapter, _ = build_adapter(stage_id=1)
    _populate_adapter_state(adapter, "req-a", "ext-a")
    _populate_adapter_state(adapter, "req-b", "ext-b")

    adapter.cleanup("req-a", "ext-a")

    assert "req-b" in adapter.upstream_exhausted_requests
    assert "req-b" in adapter.get_req_chunk
    assert "ext-b" in adapter.put_req_chunk
    assert "ext-b" in adapter.request_payload
    assert "ext-b" in adapter.code_prompt_token_ids
    assert "req-b" in adapter.request_ids_mapping


def test_cleanup_after_poll_flow(build_adapter):
    """Simulate full load_async -> poll -> finished -> cleanup cycle."""
    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
    request = _req("req-flow", RequestStatus.WAITING, external_req_id="ext-flow")

    adapter.load_async(request)

    adapter.request_ids_mapping["req-flow"] = "ext-flow"
    payload: OmniPayload = {
        "hidden_states": {"output": torch.tensor([[1.0]])},
        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
    }
    connector.get.return_value = (payload, 8)
    adapter._poll_single_request(_dequeue_load_entry(adapter, request))

    assert "req-flow" in adapter.upstream_exhausted_requests
    assert adapter.get_req_chunk["req-flow"] == 1
    assert "req-flow" in adapter.request_ids_mapping

    adapter.cleanup("req-flow", "ext-flow")

    assert "req-flow" not in adapter.upstream_exhausted_requests
    assert "req-flow" not in adapter.get_req_chunk
    assert "req-flow" not in adapter.request_ids_mapping
    assert "ext-flow" not in adapter.request_payload


def test_finish_requests_restores_status(build_adapter):
    """Abort path must pop ``requests_origin_status`` and restore pre-wait status.

    While ``process_pending_chunks`` holds a request off the scheduler queues, the
    adapter records the prior status (WAITING or RUNNING). ``finish_requests`` must
    put that status back on the live ``Request`` so base ``Scheduler.finish_requests``
    can finish bookkeeping without inconsistent state / crashes.
    """
    adapter, _ = build_adapter(stage_id=1)
    req_id = "req-abort-during-chunk"
    prior = RequestStatus.RUNNING
    request = _req(req_id, RequestStatus.WAITING_FOR_CHUNK)
    adapter.requests_origin_status[req_id] = prior
    adapter.waiting_for_chunk_running_requests.append(request)
    requests_map = {req_id: request}

    adapter.finish_requests([req_id], RequestStatus.FINISHED_ABORTED, requests_map)

    assert request.status == prior
    assert req_id not in adapter.requests_origin_status
    assert not adapter.waiting_for_chunk_running_requests


def test_finish_requests_does_not_restore_stale_status_without_connector_ownership(build_adapter):
    adapter, _ = build_adapter(stage_id=1)
    request = _req("req-restored", RequestStatus.WAITING_FOR_STREAMING_REQ)
    request.resumable = True
    adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING

    adapter.finish_requests(
        [request.request_id],
        RequestStatus.FINISHED_ABORTED,
        {request.request_id: request},
    )

    assert request.status == RequestStatus.WAITING_FOR_STREAMING_REQ
    assert request.request_id not in adapter.requests_origin_status


def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter):
    adapter, _ = build_adapter(stage_id=1)
    zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
    other = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
    adapter.waiting_for_chunk_waiting_requests = deque([zombie, other])
    adapter.waiting_for_chunk_running_requests = deque([other, zombie])
    adapter.requests_with_ready_chunks.add("req-zombie")
    adapter.upstream_exhausted_requests.add("req-zombie")
    requests_map = {
        "req-zombie": zombie,
        "req-live": other,
    }

    adapter.finish_requests(
        ["req-zombie"],
        RequestStatus.FINISHED_ABORTED,
        requests_map,
    )

    assert [req.request_id for req in adapter.waiting_for_chunk_waiting_requests] == ["req-live"]
    assert [req.request_id for req in adapter.waiting_for_chunk_running_requests] == ["req-live"]
    assert "req-zombie" not in adapter.requests_with_ready_chunks
    assert "req-zombie" not in adapter.upstream_exhausted_requests


def test_finish_requests_releases_active_stream_slot(build_adapter):
    adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=1)
    aborted = _req("req-aborted", RequestStatus.RUNNING)
    waiting = _req("req-waiting", RequestStatus.WAITING)
    waiting_queue = DummyWaitingQueue([waiting])
    running_queue: list = []
    adapter._active_streams[aborted.request_id] = aborted
    adapter._held_non_active.append(aborted)

    adapter.finish_requests(
        [aborted.request_id],
        RequestStatus.FINISHED_ABORTED,
        {aborted.request_id: aborted},
    )
    adapter.process_pending_chunks(waiting_queue, running_queue)

    assert aborted.request_id not in adapter._active_streams
    assert [request.request_id for request in adapter._held_non_active] == []
    assert list(adapter._active_streams) == [waiting.request_id]
    assert waiting.status == RequestStatus.WAITING_FOR_CHUNK


@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler])
@pytest.mark.parametrize(
    ("placement", "origin_status", "initial_status", "streaming_counter_owned"),
    [
        ("hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, False),
        ("running", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False),
        ("waiting", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False),
        ("skipped", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, True),
        ("skipped_streaming", RequestStatus.RUNNING, RequestStatus.WAITING_FOR_STREAMING_REQ, True),
        ("skipped_hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, True),
        ("skipped_non_streaming", RequestStatus.WAITING, RequestStatus.WAITING, False),
    ],
)
def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
    build_adapter,
    scheduler_cls,
    placement,
    origin_status,
    initial_status,
    streaming_counter_owned,
):
    adapter, _ = build_adapter(stage_id=1, active_stream_window=2)
    request = _req(
        "req-segment-stop",
        initial_status,
        external_req_id="ext-segment-stop",
    )
    request.resumable = True
    adapter.requests_origin_status[request.request_id] = origin_status
    if placement in {"hidden", "skipped_hidden"}:
        adapter.waiting_for_chunk_running_requests.append(request)
    adapter._active_streams[request.request_id] = request

    scheduler = scheduler_cls.__new__(scheduler_cls)
    scheduler.chunk_transfer_adapter = adapter
    scheduler.input_coordinator = None
    scheduler.requests = {request.request_id: request}
    scheduler.running = [request] if placement == "running" else []
    scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else [])
    scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
    scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)

    freed = []

    def fake_free_request(self, live, delay_free_blocks=False):
        del delay_free_blocks
        freed.append(live.request_id)
        self.requests.pop(live.request_id)
        return None, None

    scheduler._free_request = MethodType(fake_free_request, scheduler)

    first = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED)
    second = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED)

    assert len(first) == 1
    assert second == []
    assert freed == [request.request_id]
    assert request.request_id not in scheduler.requests
    assert request.request_id not in adapter._active_streams
    assert not adapter.waiting_for_chunk_running_requests
    assert scheduler.running == []
    assert list(scheduler.waiting) == []
    assert list(scheduler.skipped_waiting) == []
    assert scheduler.num_waiting_for_streaming_input == 0

    fresh_a = _req("req-fresh-a", RequestStatus.WAITING)
    fresh_b = _req("req-fresh-b", RequestStatus.WAITING)
    assert adapter._ensure_active_stream(fresh_a)
    assert adapter._ensure_active_stream(fresh_b)
    assert set(adapter._active_streams) == {fresh_a.request_id, fresh_b.request_id}


@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler])
def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapter, scheduler_cls):
    adapter, _ = build_adapter(stage_id=1, active_stream_window=2)
    request = _req("req-off-queue", RequestStatus.FINISHED_STOPPED)
    request.resumable = True
    adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING

    scheduler = scheduler_cls.__new__(scheduler_cls)
    scheduler.chunk_transfer_adapter = adapter
    scheduler.input_coordinator = None
    scheduler.requests = {request.request_id: request}
    scheduler.running = []
    scheduler.waiting = DummyWaitingQueue()
    scheduler.skipped_waiting = DummyWaitingQueue()
    scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")

    assert (
        scheduler_cls.finish_requests(
            scheduler,
            [request.request_id],
            RequestStatus.FINISHED_ABORTED,
        )
        == []
    )
    assert scheduler.requests == {request.request_id: request}
    assert request.status == RequestStatus.FINISHED_STOPPED
    assert request.request_id not in adapter.requests_origin_status


def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter):
    adapter, _ = build_adapter(stage_id=1)
    zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
    live = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
    waiting_queue = DummyWaitingQueue()
    running_queue: list = []
    adapter.waiting_for_chunk_waiting_requests = deque([zombie, live])
    adapter.waiting_for_chunk_running_requests = deque([zombie, live])

    adapter.restore_queues(
        waiting_queue,
        running_queue,
        scheduler_requests={"req-live": live},
    )

    assert [req.request_id for req in waiting_queue] == ["req-live"]
    assert [req.request_id for req in running_queue] == ["req-live"]
    assert not adapter.waiting_for_chunk_waiting_requests
    assert not adapter.waiting_for_chunk_running_requests


# ---------------------------------------------------------------
# Scheduler trigger tests
# ---------------------------------------------------------------


class _HashableRequest(SimpleNamespace):
    """SimpleNamespace that can be added to a set (needed by scheduler internals)."""

    # vLLM 0.26: update_from_output settles this counter for every scheduled
    # request; real Requests initialise it to 0 (vllm/v1/request.py).
    num_in_flight_tokens = 0

    # vLLM 0.27 (a0c092ee72): the stale-output counter that replaced
    # async_tokens_to_discard. update_from_output reads it for every scheduled
    # request, and real Requests initialise it to 0 (vllm/v1/request.py), so the
    # double needs it too.
    num_stale_output_tokens = 0

    def __hash__(self):  # type: ignore[override]  # Test requests are hashable by request ID.
        return hash(self.request_id)

    def __eq__(self, other):
        return getattr(other, "request_id", None) == self.request_id


def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: MockerFixture):
    """OmniGenerationScheduler must call adapter.cleanup when request finishes."""
    cleanup_calls = []

    adapter_mock = mocker.MagicMock()
    adapter_mock.upstream_exhausted_requests = {"req-s1"}
    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))

    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler

    scheduler = mocker.MagicMock()
    scheduler._first_chunk_express = False
    scheduler._express_min_slack_s = 0.0
    scheduler.chunk_transfer_adapter = adapter_mock
    scheduler.connector = None
    scheduler.ec_connector = None
    scheduler.perf_metrics = None
    scheduler.log_stats = False
    scheduler.recompute_kv_load_failures = False
    scheduler.structured_output_manager = mocker.MagicMock()
    scheduler.structured_output_manager.accept_tokens.return_value = True
    scheduler.finished_req_ids_dict = {}
    scheduler.kv_cache_manager.take_events.return_value = None
    scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
    scheduler.kv_event_publisher = mocker.MagicMock()

    request = _HashableRequest(
        request_id="req-s1",
        external_req_id="ext-s1",
        status=RequestStatus.RUNNING,
        is_finished=lambda: False,
        num_computed_tokens=10,
        num_prompt_tokens=10,
        prompt_token_ids=list(range(10)),
        num_output_placeholders=0,
        sampling_params=None,
        pooling_params=None,
        stop_reason=None,
        client_index=0,
        take_events=lambda: [],
        trace_headers=None,
        has_encoder_inputs=False,
        take_prefill_stats=lambda: None,
        num_nans_in_logits=0,
        get_finished_reason=lambda: "stop",
    )
    scheduler.requests = {"req-s1": request}

    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
    scheduler._free_request = mocker.MagicMock(return_value=(None, None))
    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
    scheduler.running = [request]
    scheduler.waiting = mocker.MagicMock()
    scheduler.waiting.remove_requests = mocker.MagicMock()
    scheduler.make_stats = mocker.MagicMock(return_value=None)

    scheduler_output = SimpleNamespace(
        num_scheduled_tokens={"req-s1": 10},
        scheduled_spec_decode_tokens={},
        num_invalid_spec_tokens=0,
    )
    model_runner_output = SimpleNamespace(
        sampled_token_ids=None,
        logprobs=None,
        prompt_logprobs_dict={},
        pooler_output=None,
        num_nans_in_logits=None,
        kv_connector_output=None,
        cudagraph_stats=None,
        req_id_to_index={"req-s1": 0},
        routed_experts=None,
        routed_experts_dict=None,
    )

    OmniGenerationScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)

    assert len(cleanup_calls) == 1
    args, _ = cleanup_calls[0]
    assert args[0] == "req-s1"
    assert args[1] == "ext-s1"


def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerFixture):
    """OmniARScheduler should enqueue save; adapter cleanup is handled in save thread."""
    cleanup_calls = []
    save_calls = []

    adapter_mock = mocker.MagicMock()
    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
    adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))

    scheduler = mocker.MagicMock()
    scheduler.chunk_transfer_adapter = adapter_mock
    bind_omits_transfer_helpers(scheduler)
    scheduler.connector = None
    scheduler.perf_metrics = None
    scheduler.log_stats = False
    scheduler.recompute_kv_load_failures = False
    scheduler.structured_output_manager = mocker.MagicMock()
    scheduler.structured_output_manager.accept_tokens.return_value = True
    scheduler.finished_req_ids_dict = {}
    scheduler.kv_cache_manager = mocker.MagicMock()
    scheduler.kv_cache_manager.take_events.return_value = None
    scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
    scheduler.kv_event_publisher = mocker.MagicMock()
    scheduler.waiting_for_transfer_free = set()
    scheduler.transfer_triggered_requests = set()
    scheduler.active_kv_transfers = set()

    request = _HashableRequest(
        request_id="req-ar",
        external_req_id="ext-ar",
        status=RequestStatus.RUNNING,
        is_finished=lambda: False,
        num_computed_tokens=1,
        num_prompt_tokens=1,
        prompt_token_ids=[1],
        num_output_placeholders=0,
        sampling_params=None,
        pooling_params=None,
        stop_reason=None,
        client_index=0,
        take_events=lambda: [],
        trace_headers=None,
        has_encoder_inputs=False,
        take_prefill_stats=lambda: None,
        num_nans_in_logits=0,
        get_finished_reason=lambda: "stop",
    )
    scheduler.requests = {"req-ar": request}

    scheduler._update_request_with_output = mocker.MagicMock(return_value=([], True))
    scheduler._process_kv_transfer_trigger = mocker.MagicMock(return_value=False)
    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
    scheduler._free_request = mocker.MagicMock(return_value=(None, None))
    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
    scheduler.running = [request]
    scheduler.waiting = mocker.MagicMock()
    scheduler.waiting.remove_requests = mocker.MagicMock()
    scheduler.make_spec_decoding_stats = mocker.MagicMock(return_value=None)
    scheduler.make_stats = mocker.MagicMock(return_value=None)

    scheduler_output = SimpleNamespace(
        num_scheduled_tokens={"req-ar": 1},
        scheduled_spec_decode_tokens={},
        num_invalid_spec_tokens=0,
    )
    model_runner_output = SimpleNamespace(
        sampled_token_ids=[[123]],
        logprobs=None,
        prompt_logprobs_dict={},
        pooler_output=None,
        num_nans_in_logits=None,
        kv_connector_output=None,
        cudagraph_stats=None,
        req_id_to_index={"req-ar": 0},
        kv_extracted_req_ids=None,
        routed_experts=None,
        routed_experts_dict=None,
    )

    OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)

    assert len(cleanup_calls) == 0
    assert len(save_calls) == 1


def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
    """``OmniARScheduler.finish_requests`` must run chunk adapter hook before vLLM base."""
    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler

    order: list[str] = []

    adapter = mocker.MagicMock()

    def _adapter_finish(request_ids, finished_status, requests):
        order.append("adapter")
        return []

    adapter.finish_requests.side_effect = _adapter_finish

    def _super_finish(_self, request_ids, finished_status):
        order.append("super")
        return []

    sched = OmniARScheduler.__new__(OmniARScheduler)
    sched.chunk_transfer_adapter = adapter
    sched.requests = {}
    sched.running = []
    sched.waiting = []

    with patch.object(VLLMScheduler, "finish_requests", _super_finish):
        OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)

    assert order == ["adapter", "super"]


def _parked_sender_scheduler(adapter, session):
    """A minimal OmniARScheduler holding one parked streaming session."""
    scheduler = OmniARScheduler.__new__(OmniARScheduler)
    scheduler.chunk_transfer_adapter = adapter
    scheduler.input_coordinator = None
    scheduler.requests = {session.request_id: session}
    scheduler.running = []
    scheduler.waiting = DummyWaitingQueue()
    scheduler.skipped_waiting = DummyWaitingQueue([session])
    scheduler.num_waiting_for_streaming_input = 1
    scheduler.finished_req_ids_dict = {}

    def _free_request(self, request, delay_free_blocks=False):
        del delay_free_blocks
        self.requests.pop(request.request_id, None)
        return None, None

    scheduler._free_request = MethodType(_free_request, scheduler)
    return scheduler


@pytest.mark.parametrize("send_first_chunk", [False, True])
def test_parked_streaming_final_update_emits_downstream_terminal(build_adapter, send_first_chunk):
    """#6670: a final update on a parked sender must terminate the receiver.

    Upstream turns "final update while parked in WAITING_FOR_STREAMING_REQ"
    into a silent local abort. The downstream stage only learns a stream ended
    from a chunk carrying ``finished``, so the abort used to leave it in
    WAITING_FOR_CHUNK until the 600s input deadline.
    """
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    session = _req(
        "req-parked-final",
        RequestStatus.WAITING_FOR_STREAMING_REQ,
        external_req_id="ext-parked-final",
    )
    session.resumable = True
    session._omni_segment_generation = 0
    # A segment stop queues its boundary before the final update arrives.
    # The downstream request is already prewarmed even if the sender has not
    # dequeued this first task, so it still needs a terminal in both cases.
    adapter.save_async(None, session, is_segment_finished=True)
    assert session.external_req_id not in adapter.put_req_chunk
    if send_first_chunk:
        adapter._send_single_request(adapter._pending_save_reqs.popleft())
        connector.put.reset_mock()

    scheduler = _parked_sender_scheduler(adapter, session)

    final_update = _req(session.request_id, RequestStatus.WAITING)
    final_update.resumable = False
    OmniARScheduler.add_request(scheduler, final_update)

    # Local teardown happened...
    assert session.request_id not in scheduler.requests
    # ...but only after a terminal chunk was queued for the next stage.
    assert session.status == RequestStatus.FINISHED_STOPPED
    assert any(task["is_finished"] for task in adapter._pending_save_reqs)
    while adapter._pending_save_reqs:
        adapter._send_single_request(adapter._pending_save_reqs.popleft())

    connector.put.assert_called_once()
    assert connector.put.call_args.kwargs["put_key"] == f"ext-parked-final_1_{int(send_first_chunk)}"
    payload = connector.put.call_args.kwargs["data"]
    assert bool(payload.meta.finished.item()) is True
    assert bool(payload.meta.is_segment_finished.item()) is False


def test_queued_terminal_chunk_survives_sender_cleanup(build_adapter):
    """#6670: ``cleanup_sender`` must not discard an already-queued terminal."""
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal")
    request.resumable = False

    adapter.save_async(None, request, is_segment_finished=False)
    assert len(adapter._pending_save_reqs) == 1

    # The scheduler tears the request down before the save_loop drains.
    adapter.cleanup_sender("ext-terminal")

    adapter._send_single_request(adapter._pending_save_reqs.popleft())
    connector.put.assert_called_once()
    assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True
    # The terminal put reclaims the generation it kept alive.
    assert "ext-terminal" not in adapter._sender_tokens


def test_non_terminal_chunk_is_still_dropped_by_sender_cleanup(build_adapter):
    """The #6670 fence is terminal-only: ordinary queued chunks still drop."""
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    request = _req("req-mid", RequestStatus.RUNNING, external_req_id="ext-mid")
    request.resumable = True

    adapter.save_async(None, request, is_segment_finished=False)
    adapter.cleanup_sender("ext-mid")

    adapter._send_single_request(adapter._pending_save_reqs.popleft())
    connector.put.assert_not_called()


def test_final_update_on_non_sender_stage_keeps_upstream_abort(build_adapter):
    """A final stage never put a chunk, so it must keep the upstream path."""
    adapter, _connector = build_adapter(stage_id=2, model_mode="generation")
    session = _req(
        "req-final-stage",
        RequestStatus.WAITING_FOR_STREAMING_REQ,
        external_req_id="ext-final-stage",
    )
    session.resumable = True
    scheduler = _parked_sender_scheduler(adapter, session)

    final_update = _req(session.request_id, RequestStatus.WAITING)
    final_update.resumable = False

    with patch.object(VLLMScheduler, "add_request") as base_add:
        OmniARScheduler.add_request(scheduler, final_update)

    base_add.assert_called_once()
    assert not adapter._pending_save_reqs


def _queue_terminal(
    adapter,
    external_req_id: str,
    req_id: str = "req-terminal",
    segment_generation: int = 0,
) -> None:
    """Queue one request-terminal chunk on the adapter's save queue.

    ``segment_generation`` must match the adapter's dedup watermark, which a
    preceding segment stop advances; otherwise ``save_async`` drops the chunk
    as a late duplicate.
    """
    request = _req(req_id, RequestStatus.FINISHED_STOPPED, external_req_id=external_req_id)
    request.resumable = False
    request._omni_segment_generation = segment_generation
    adapter.save_async(None, request, is_segment_finished=False)


def _assert_terminal_put(connector) -> None:
    connector.put.assert_called_once()
    assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True


def test_terminal_chunk_survives_cleanup_after_dequeue(build_adapter):
    """#6670: the fence must hold across the connector handoff, not just the queue.

    ``_finish_parked_streaming_session`` queues the terminal and then calls
    ``finish_requests``, so ``cleanup_sender`` can land *after* the save_loop
    has already dequeued the terminal. Releasing the fence at dequeue left
    that interleaving cancelling the chunk before ``put`` -- the same lost
    terminal, and the same downstream WAITING_FOR_CHUNK hang.
    """
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    _queue_terminal(adapter, "ext-late-cleanup")
    task = adapter._pending_save_reqs.popleft()

    real_send = adapter._send_single_request_for_generation

    def send_with_cleanup_in_flight(inner_task, sender_token=None):
        # The scheduler thread reaches cleanup_sender here: the save_loop owns
        # the terminal but has not handed it to the connector yet.
        adapter.cleanup_sender("ext-late-cleanup")
        return real_send(inner_task, sender_token)

    adapter._send_single_request_for_generation = send_with_cleanup_in_flight
    adapter._send_single_request(task)

    _assert_terminal_put(connector)
    # The generation the fence kept alive is reclaimed, not leaked.
    assert "ext-late-cleanup" not in adapter._sender_tokens
    assert "ext-late-cleanup" not in adapter.put_req_chunk


def test_terminal_chunk_survives_cleanup_racing_an_inflight_sibling(build_adapter):
    """#6670: an in-flight ordinary chunk must not reclaim a fenced generation.

    ``_send_single_request``'s ``finally`` also tears a cancelled generation
    down. If the stop and the abort both land while an earlier chunk is on the
    wire, that teardown would unregister the token and the queued terminal
    would then fail the identity check.
    """
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    running = _req("req-sibling", RequestStatus.RUNNING, external_req_id="ext-sibling")
    running.resumable = True
    adapter.save_async(None, running, is_segment_finished=True)
    ordinary_task = adapter._pending_save_reqs.popleft()

    real_send = adapter._send_single_request_for_generation

    def send_then_stop_and_abort(inner_task, sender_token=None):
        result = real_send(inner_task, sender_token)
        # Terminal stop and abort both land while this chunk is on the wire.
        # The segment stop above already armed the next segment's watermark.
        _queue_terminal(adapter, "ext-sibling", req_id="req-sibling", segment_generation=1)
        adapter.cleanup_sender("ext-sibling")
        return result

    adapter._send_single_request_for_generation = send_then_stop_and_abort
    adapter._send_single_request(ordinary_task)
    adapter._send_single_request_for_generation = real_send

    assert len(adapter._pending_save_reqs) == 1
    connector.put.reset_mock()
    adapter._send_single_request(adapter._pending_save_reqs.popleft())

    _assert_terminal_put(connector)
    assert "ext-sibling" not in adapter._sender_tokens


@pytest.mark.parametrize("finish_inside_put", [False, True])
def test_successful_put_keeps_terminal_chunk_sequence(build_adapter, finish_inside_put):
    """A cleanup during put must not reuse a key already consumed downstream."""
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    receiver, receiver_connector = build_adapter(stage_id=2, model_mode="generation")
    session = _req("req-sequence", RequestStatus.WAITING_FOR_STREAMING_REQ, external_req_id="ext-sequence")
    session.resumable = True
    session._omni_segment_generation = 0
    scheduler = _parked_sender_scheduler(adapter, session)
    final_update = _req(session.request_id, RequestStatus.WAITING)
    final_update.resumable = False
    downstream = _req("req-downstream", RequestStatus.WAITING, external_req_id=session.external_req_id)
    receiver.request_ids_mapping[downstream.request_id] = session.external_req_id
    payloads = {}

    def receive_next():
        entry = _dequeue_load_entry(receiver, downstream)
        return receiver._poll_single_request(entry)

    def put_then_finish(**kwargs):
        payload = kwargs["data"]
        payloads[kwargs["put_key"]] = {
            "meta": {
                "finished": bool(payload.meta.finished.item()),
                "is_segment_finished": bool(payload.meta.is_segment_finished.item()),
            }
        }
        if not payload.meta.finished.item():
            # Consume the ordinary boundary while put is still in flight.
            assert receive_next()
            if finish_inside_put:
                scheduler.add_request(final_update)
        return True, 1, {}

    def get_by_key(_from_stage, _to_stage, key, metadata=None):
        assert metadata is None
        payload = payloads.pop(key, None)
        return (payload, 1) if payload is not None else None

    connector.put.side_effect = put_then_finish
    receiver_connector.get.side_effect = get_by_key
    adapter.save_async(None, session, is_segment_finished=True)
    adapter._send_single_request(adapter._pending_save_reqs.popleft())
    if not finish_inside_put:
        scheduler.add_request(final_update)
    while adapter._pending_save_reqs:
        adapter._send_single_request(adapter._pending_save_reqs.popleft())

    assert [call.kwargs["put_key"] for call in connector.put.call_args_list] == ["ext-sequence_1_0", "ext-sequence_1_1"]
    assert receive_next()
    assert downstream.request_id in receiver.upstream_exhausted_requests
    assert not payloads
    assert session.external_req_id not in adapter._sender_tokens


def test_late_successful_put_cannot_advance_replacement_generation(build_adapter):
    """The post-put guard must retain identity checks while relaxing cancellation."""
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    request = _req("req-old", RequestStatus.RUNNING, external_req_id="ext-replaced")
    request.resumable = True
    adapter.save_async(None, request, is_segment_finished=True)
    task = adapter._pending_save_reqs.popleft()
    old_token = task["sender_token"]
    replacement_token = type(old_token)()

    def replace_during_put(**kwargs):
        # Model a completion that no longer owns its external id. The old
        # terminal fence must never authorize writes to the successor's state.
        with adapter._sender_state_lock:
            old_token.cancelled = True
            old_token.terminal_pending = True
            adapter._sender_tokens[request.external_req_id] = replacement_token
            adapter.put_req_chunk[request.external_req_id] = 4
            adapter.ramp_chunk_count[request.external_req_id] = 2
        return True, 1, {}

    connector.put.side_effect = replace_during_put
    adapter._send_single_request(task)
    assert adapter._sender_tokens[request.external_req_id] is replacement_token
    assert adapter.put_req_chunk[request.external_req_id] == 4
    assert adapter.ramp_chunk_count[request.external_req_id] == 2


def test_new_request_cannot_join_a_generation_with_a_queued_terminal(build_adapter):
    """#6670: the fence must not readmit a request reusing the external id.

    ``save_async`` rejects chunks whose generation is ``cancelled``; that is
    what stops a new request from inheriting a retiring generation's token and
    chunk counters. Deferring the reclaim must not also defer that retirement.
    """
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    _queue_terminal(adapter, "ext-shared", req_id="req-old")
    adapter.cleanup_sender("ext-shared")

    newcomer = _req("req-new", RequestStatus.RUNNING, external_req_id="ext-shared")
    newcomer.resumable = True
    adapter.save_async(None, newcomer, is_segment_finished=True)

    # Rejected loudly, so the scheduler fails the newcomer now instead of its
    # chunks being silently discarded at send time.
    assert len(adapter._pending_save_reqs) == 1
    assert adapter.collect_failed_send_request_ids() == {"req-new": "previous sender generation is still draining"}

    adapter._send_single_request(adapter._pending_save_reqs.popleft())
    _assert_terminal_put(connector)
    assert "ext-shared" not in adapter._sender_tokens


def test_failed_terminal_put_runs_the_deferred_cleanup(build_adapter):
    """#6670: a silent connector drop must not strand the deferred cleanup.

    ``connector.put`` returning False skips the success path that would have
    called ``cleanup``. The reclaim ``cleanup_sender`` handed over is owed on
    that path too, or the whole per-request sender state leaks for good.
    """
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    _queue_terminal(adapter, "ext-put-failed")
    adapter.cleanup_sender("ext-put-failed")
    connector.put.return_value = (False, 0, {})

    adapter._send_single_request(adapter._pending_save_reqs.popleft())

    assert "ext-put-failed" not in adapter._sender_tokens
    assert "ext-put-failed" not in adapter.put_req_chunk


def test_raising_terminal_put_runs_the_deferred_cleanup(build_adapter):
    """#6670: the same reclaim is owed when the connector raises."""
    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
    _queue_terminal(adapter, "ext-put-raised")
    adapter.cleanup_sender("ext-put-raised")
    connector.put.side_effect = RuntimeError("connector down")

    # ``save_loop`` catches this and records a send failure.
    with pytest.raises(RuntimeError):
        adapter._send_single_request(adapter._pending_save_reqs.popleft())

    assert "ext-put-raised" not in adapter._sender_tokens
    assert "ext-put-raised" not in adapter.put_req_chunk


def test_wire_round_trip_struct_to_dict_contract():
    """Pin the wire contract: encoding ``OmniPayloadStruct`` and decoding it
    yields a dict equivalent to ``to_dict(struct)``.

    The chunk-adapter sender uses struct attribute access while the receiver
    uses dict-key access. This works only because ``OmniMsgpackDecoder`` has
    no target type and decodes structs back to plain dicts. If this test
    breaks, the receiver's dict access will silently drop fields or KeyError.
    """
    from vllm_omni.data_entry_keys import CodesStruct, to_dict
    from vllm_omni.distributed.omni_connectors.utils.serialization import (
        OmniMsgpackDecoder,
        OmniMsgpackEncoder,
    )

    struct = OmniPayloadStruct(
        meta=MetaStruct(
            finished=torch.tensor(True, dtype=torch.bool),
            left_context_size=12,
        ),
        codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.int64)),
    )

    encoded = OmniMsgpackEncoder().encode(struct)
    decoded = OmniMsgpackDecoder().decode(encoded)

    assert isinstance(decoded, dict)
    assert isinstance(decoded["meta"], dict)
    assert isinstance(decoded["meta"]["finished"], torch.Tensor)
    assert bool(decoded["meta"]["finished"].item()) is True
    assert decoded["meta"]["left_context_size"] == 12
    assert torch.equal(decoded["codes"]["audio"], torch.tensor([1, 2, 3], dtype=torch.int64))

    expected = to_dict(struct)
    assert set(decoded.keys()) == set(expected.keys())
    assert set(decoded["meta"].keys()) == set(expected["meta"].keys())
    assert set(decoded["codes"].keys()) == set(expected["codes"].keys())


# ---------------------------------------------------------------
# Deferred finish for upstream-completed requests
# ---------------------------------------------------------------


def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
    """Build a mock scheduler with requests queued for deferred finish."""
    adapter_mock = mocker.MagicMock()
    adapter_mock.upstream_exhausted_requests = {r.request_id for r in pending_finish_reqs}
    cleanup_calls = []
    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))

    scheduler = mocker.MagicMock()
    scheduler.chunk_transfer_adapter = adapter_mock
    scheduler.connector = None
    scheduler.ec_connector = None
    scheduler.perf_metrics = None
    scheduler.log_stats = False
    scheduler.recompute_kv_load_failures = False
    scheduler.structured_output_manager = mocker.MagicMock()
    scheduler.structured_output_manager.accept_tokens.return_value = True
    scheduler.finished_req_ids_dict = {}
    scheduler.kv_cache_manager.take_events.return_value = None
    scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
    scheduler.kv_event_publisher = mocker.MagicMock()
    scheduler._pending_finish_reqs = list(pending_finish_reqs)

    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
    scheduler._free_request = mocker.MagicMock(return_value=(None, None))
    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
    scheduler.running = list(running)
    scheduler.waiting = mocker.MagicMock()
    scheduler.waiting.remove_requests = mocker.MagicMock()
    scheduler.make_stats = mocker.MagicMock(return_value=None)
    scheduler.requests = {r.request_id: r for r in running}

    scheduler_output = SimpleNamespace(
        num_scheduled_tokens={},
        scheduled_spec_decode_tokens={},
        num_invalid_spec_tokens=0,
    )
    model_runner_output = SimpleNamespace(
        sampled_token_ids=None,
        logprobs=None,
        prompt_logprobs_dict={},
        pooler_output=None,
        num_nans_in_logits=None,
        kv_connector_output=None,
        cudagraph_stats=None,
        req_id_to_index={},
    )
    return scheduler, scheduler_output, model_runner_output, cleanup_calls


def test_deferred_finish_emits_finished_output(mocker: MockerFixture):
    """A request whose upstream completed with no remaining tokens should
    emit a FINISHED EngineCoreOutput, free resources, and clean up adapter state."""
    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler

    request = _HashableRequest(
        request_id="req-df1",
        external_req_id="ext-df1",
        status=RequestStatus.RUNNING,
        is_finished=lambda: False,
        num_computed_tokens=16,
        num_prompt_tokens=16,
        prompt_token_ids=list(range(1, 17)),
        num_output_placeholders=0,
        sampling_params=None,
        pooling_params=None,
        stop_reason=None,
        client_index=0,
        take_events=lambda: [],
        trace_headers=None,
        has_encoder_inputs=False,
        take_prefill_stats=lambda: None,
        num_nans_in_logits=0,
        get_finished_reason=lambda: "stop",
    )
    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
        mocker,
        running=[request],
        pending_finish_reqs=[request],
    )
    scheduler._free_request.return_value = ({"mock": "kv_params"}, None)

    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)

    assert request.status == RequestStatus.FINISHED_STOPPED
    scheduler._handle_stopped_request.assert_called_once_with(request)
    scheduler._free_request.assert_called_once_with(request)
    assert len(cleanup_calls) == 1
    assert cleanup_calls[0][0] == ("req-df1", "ext-df1")

    eco = result[0]
    assert len(eco.outputs) == 1
    assert eco.outputs[0].request_id == "req-df1"
    assert eco.outputs[0].finish_reason == "stop"
    assert eco.outputs[0].kv_transfer_params == {"mock": "kv_params"}
    assert scheduler._pending_finish_reqs == []


def test_deferred_finish_empty_prompt(mocker: MockerFixture):
    """A request that never received any tokens (finished immediately upstream)
    should still emit a FINISHED output and clean up."""
    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler

    request = _HashableRequest(
        request_id="req-df2",
        external_req_id="ext-df2",
        status=RequestStatus.WAITING,
        is_finished=lambda: False,
        num_computed_tokens=0,
        num_prompt_tokens=0,
        prompt_token_ids=[],
        num_output_placeholders=0,
        sampling_params=None,
        pooling_params=None,
        stop_reason=None,
        client_index=0,
        take_events=lambda: [],
        trace_headers=None,
        has_encoder_inputs=False,
        take_prefill_stats=lambda: None,
        num_nans_in_logits=0,
        get_finished_reason=lambda: "stop",
    )
    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
        mocker,
        running=[],
        pending_finish_reqs=[request],
    )

    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)

    assert request.status == RequestStatus.FINISHED_STOPPED
    scheduler._free_request.assert_called_once_with(request)
    assert len(cleanup_calls) == 1
    eco = result[0]
    assert len(eco.outputs) == 1
    assert eco.outputs[0].finish_reason == "stop"
    assert scheduler._pending_finish_reqs == []


def test_deferred_finish_skips_already_finished(mocker: MockerFixture):
    """A request that was aborted between schedule() and update_from_output()
    should be skipped without emitting output or freeing resources twice."""
    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler

    request = _HashableRequest(
        request_id="req-df3",
        external_req_id="ext-df3",
        status=RequestStatus.FINISHED_ABORTED,
        is_finished=lambda: True,
        num_computed_tokens=0,
        num_prompt_tokens=0,
        prompt_token_ids=[],
        num_output_placeholders=0,
        sampling_params=None,
        pooling_params=None,
        stop_reason=None,
        client_index=0,
        take_events=lambda: [],
        trace_headers=None,
        has_encoder_inputs=False,
        take_prefill_stats=lambda: None,
        num_nans_in_logits=0,
        get_finished_reason=lambda: "stop",
    )
    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
        mocker,
        running=[],
        pending_finish_reqs=[request],
    )

    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)

    scheduler._handle_stopped_request.assert_not_called()
    scheduler._free_request.assert_not_called()
    assert len(cleanup_calls) == 0
    assert 0 not in result
    assert scheduler._pending_finish_reqs == []


def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture):
    """When _handle_stopped_request returns False (resumable request), the
    output must still be emitted so the client stream doesn't hang."""
    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler

    request = _HashableRequest(
        request_id="req-df4",
        external_req_id="ext-df4",
        status=RequestStatus.RUNNING,
        is_finished=lambda: False,
        num_computed_tokens=16,
        num_prompt_tokens=16,
        prompt_token_ids=list(range(1, 17)),
        num_output_placeholders=0,
        sampling_params=None,
        pooling_params=None,
        stop_reason=None,
        client_index=0,
        take_events=lambda: [],
        trace_headers=None,
        has_encoder_inputs=False,
        take_prefill_stats=lambda: None,
        num_nans_in_logits=0,
        get_finished_reason=lambda: "stop",
    )
    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
        mocker,
        running=[request],
        pending_finish_reqs=[request],
    )
    scheduler._handle_stopped_request.return_value = False

    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)

    scheduler._handle_stopped_request.assert_called_once_with(request)
    scheduler._free_request.assert_not_called()
    assert len(cleanup_calls) == 0

    eco = result[0]
    assert len(eco.outputs) == 1
    assert eco.outputs[0].request_id == "req-df4"
    assert eco.outputs[0].finish_reason == "stop"
    assert eco.outputs[0].kv_transfer_params is None
    assert scheduler._pending_finish_reqs == []


# ---------------------------------------------------------------
# Zombie purge tests (regression for vllm-project/vllm-omni#3736)
# ---------------------------------------------------------------


def test_process_pending_chunks_purges_zombies_in_running_deque(
    build_adapter,
):
    """A request aborted while parked in ``waiting_for_chunk_running_requests``
    is no longer tracked by ``scheduler.requests`` once
    ``Scheduler._free_request`` runs. ``process_pending_chunks`` must drop
    that zombie *before* ``restore_queues`` would re-inject it onto the
    running queue.

    Regression for https://github.com/vllm-project/vllm-omni/issues/3736
    (engine-core ``KeyError`` on aborts under chunk-transfer pressure).
    """
    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)

    live_req = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
    zombie_req = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
    adapter.waiting_for_chunk_running_requests.append(live_req)
    adapter.waiting_for_chunk_running_requests.append(zombie_req)
    # Mirror state the adapter would carry for an in-flight request so we can
    # later assert ``cleanup_receiver`` actually fired against the zombie.
    adapter.requests_origin_status[zombie_req.request_id] = RequestStatus.RUNNING
    adapter.requests_origin_status[live_req.request_id] = RequestStatus.RUNNING

    waiting_queue = DummyWaitingQueue()
    running_queue: list = []
    # Simulate the post-abort scheduler: only ``live-1`` is still tracked.
    scheduler_requests = {live_req.request_id: live_req}

    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)

    # 1. Zombie was popped from the deque.
    assert zombie_req not in adapter.waiting_for_chunk_running_requests
    # 2. Live request is still in the deque.
    assert live_req in adapter.waiting_for_chunk_running_requests
    # 3. ``cleanup_receiver`` ran for the zombie (drops origin-status mapping
    #    and unregisters it so a late load/poll is dropped too).
    assert zombie_req.request_id not in adapter.requests_origin_status
    assert zombie_req.request_id not in adapter._registered_load_entries
    # 4. Live request's bookkeeping is untouched.
    assert adapter.requests_origin_status[live_req.request_id] == RequestStatus.RUNNING

    # 5. ``restore_queues`` (which the scheduler runs in its ``finally``
    #    clause) now only re-injects the live request -- the zombie is
    #    gone, so the worker's ``_update_states`` cannot crash on it.
    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
    assert running_queue == [live_req]


def test_process_pending_chunks_purges_zombies_in_waiting_deque(build_adapter):
    """Zombie purge applies symmetrically to the waiting-side deque.

    Regression for https://github.com/vllm-project/vllm-omni/issues/3736 --
    aborted requests that landed in ``waiting_for_chunk_waiting_requests``
    must also be removed before ``restore_queues`` re-injects them, since
    ``restore_queues`` uses ``waiting_queue.add_request`` rather than
    ``running_queue.extend`` for that path.
    """
    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)

    live_req = _req("live-w", RequestStatus.WAITING_FOR_CHUNK)
    zombie_req = _req("zombie-w", RequestStatus.WAITING_FOR_CHUNK)
    adapter.waiting_for_chunk_waiting_requests.append(live_req)
    adapter.waiting_for_chunk_waiting_requests.append(zombie_req)

    waiting_queue = DummyWaitingQueue()
    running_queue: list = []
    scheduler_requests = {live_req.request_id: live_req}

    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)

    assert live_req in adapter.waiting_for_chunk_waiting_requests
    assert zombie_req not in adapter.waiting_for_chunk_waiting_requests
    assert zombie_req.request_id not in adapter._registered_load_entries

    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
    assert waiting_queue == [live_req]


def test_purge_preserves_live_order_with_interleaved_zombies(build_adapter):
    """Interleaved live/zombie ordering -- pin the ``popleft`` + append-live
    in-place filter at chunk_transfer_adapter._purge_untracked_chunk_requests:
    survivor order must match insertion order, and ``cleanup_receiver`` must
    run exactly once per zombie.
    """
    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)

    live1 = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
    zombie1 = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
    live2 = _req("live-2", RequestStatus.WAITING_FOR_CHUNK)
    zombie2 = _req("zombie-2", RequestStatus.WAITING_FOR_CHUNK)
    for req in (live1, zombie1, live2, zombie2):
        adapter.waiting_for_chunk_running_requests.append(req)

    waiting_queue = DummyWaitingQueue()
    running_queue: list = []
    scheduler_requests = {live1.request_id: live1, live2.request_id: live2}

    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)

    assert list(adapter.waiting_for_chunk_running_requests) == [live1, live2]
    assert zombie1.request_id not in adapter._registered_load_entries
    assert zombie2.request_id not in adapter._registered_load_entries


def test_restore_queues_purges_late_aborts_after_process_pending_chunks(
    build_adapter,
):
    """Race window guard: an abort can fire between ``process_pending_chunks``
    and the scheduler's ``finally``-clause ``restore_queues`` call. The
    second purge inside ``restore_queues`` must drop the now-untracked
    request so it does not get re-injected onto ``running_queue``.
    """
    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)

    req = _req("late-abort", RequestStatus.WAITING_FOR_CHUNK)
    adapter.waiting_for_chunk_running_requests.append(req)

    waiting_queue = DummyWaitingQueue()
    running_queue: list = []

    # Tick N: process_pending_chunks runs while req is still tracked.
    scheduler_requests: dict = {req.request_id: req}
    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
    assert req in adapter.waiting_for_chunk_running_requests

    # Mid-tick: abort fires and the scheduler's free path deletes the entry.
    del scheduler_requests[req.request_id]

    # finally: restore_queues sees the now-untracked req and must drop it
    # instead of blindly extending it onto running_queue.
    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
    assert running_queue == []
    assert req.request_id not in adapter._registered_load_entries


def test_purge_is_noop_on_empty_deques(build_adapter):
    """Empty deques short-circuit -- guards against any accidental
    ``IndexError`` from ``popleft`` on an empty deque if a future caller
    shadows the empty check.
    """
    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
    assert len(adapter.waiting_for_chunk_waiting_requests) == 0
    assert len(adapter.waiting_for_chunk_running_requests) == 0

    waiting_queue = DummyWaitingQueue()
    running_queue: list = []

    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests={})
    assert len(adapter.waiting_for_chunk_waiting_requests) == 0
    assert len(adapter.waiting_for_chunk_running_requests) == 0
    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
    assert running_queue == []
    assert waiting_queue == []


# --------------------------------------------------------------------------- #
#  Chunk-wait deadline (RFC #4855 R1.1, issue #3833)
#
#  Before this, the async-chunk path had no deadline of any kind, so a dropped
#  terminal chunk or an upstream stage that died mid-stream parked the request
#  in WAITING_FOR_CHUNK forever. None of these scenarios were covered.
# --------------------------------------------------------------------------- #


def _park_in_chunk_wait(adapter, request, *, waiting=True):
    """Drive a request through one process_pending_chunks round into the wait."""
    queue = DummyWaitingQueue([request]) if waiting else [request]
    if waiting:
        adapter.process_pending_chunks(queue, [])
    else:
        adapter.process_pending_chunks(DummyWaitingQueue(), queue)
    return queue


def test_chunk_wait_clock_starts_when_a_request_parks(build_adapter):
    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
    request = _req("r1", RequestStatus.WAITING)

    _park_in_chunk_wait(adapter, request)

    assert request.status == RequestStatus.WAITING_FOR_CHUNK
    assert "r1" in adapter._waiting_since


def test_chunk_wait_does_not_expire_before_the_deadline(build_adapter):
    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
    _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))

    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
    assert "r1" in adapter._waiting_since


def test_dropped_terminal_chunk_expires_the_request(build_adapter):
    """The #3833 scenario: upstream stops sending and never marks the stream done."""
    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
    _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))

    adapter._waiting_since["r1"] -= 601.0

    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"r1"}
    # Cleared, so a second sweep does not re-report the same request.
    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()


def test_arriving_chunk_resets_the_clock(build_adapter):
    """A slow but healthy stream must never expire: the deadline measures stall
    time between chunks, not the lifetime of the stream."""
    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
    request = _req("r1", RequestStatus.WAITING)
    queue = _park_in_chunk_wait(adapter, request)
    adapter.restore_queues(queue, [], scheduler_requests={"r1": request})

    # Age the wait to just under the deadline, then deliver a chunk.
    adapter._waiting_since["r1"] -= 599.0
    adapter._finished_load_reqs.add("r1")
    adapter.process_pending_chunks(queue, [])

    assert "r1" not in adapter._waiting_since
    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()


def test_a_disabled_deadline_never_expires(build_adapter):
    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
    _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
    adapter._waiting_since["r1"] -= 10_000.0

    assert adapter.collect_timed_out_request_ids(timeout_s=0.0) == set()
    assert adapter.collect_timed_out_request_ids(timeout_s=-1.0) == set()


def test_finished_request_leaves_no_stale_timestamp(build_adapter):
    """Otherwise a completed request would be reported as timed out later."""
    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
    request = _req("r1", RequestStatus.WAITING)
    _park_in_chunk_wait(adapter, request)

    adapter.finish_requests(["r1"], RequestStatus.FINISHED_STOPPED, {"r1": request})

    assert "r1" not in adapter._waiting_since
    assert adapter.collect_timed_out_request_ids(timeout_s=0.001) == set()


def test_aborted_request_leaves_no_stale_timestamp(build_adapter):
    adapter, _ = build_adapter(stage_id=1, model_mode="ar")
    _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))

    adapter.cleanup_receiver("r1")

    assert "r1" not in adapter._waiting_since


def test_expiry_is_per_request(build_adapter):
    """One stalled stream must not take down its healthy neighbours."""
    adapter, _ = build_adapter(stage_id=1, model_mode="ar", max_num_seqs=4)
    stalled = _req("stalled", RequestStatus.WAITING)
    healthy = _req("healthy", RequestStatus.WAITING)
    adapter.process_pending_chunks(DummyWaitingQueue([stalled, healthy]), [])

    adapter._waiting_since["stalled"] -= 601.0

    assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
    assert "healthy" in adapter._waiting_since


def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter):
    from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
        talker2code2wav_async_chunk,
    )

    adapter, _ = build_adapter(stage_id=1, connector_extra={"codec_chunk_frames": 1})
    external_req_id = "ext-native-abort"
    first = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
    first.resumable = True
    first.model_intermediate_buffer = None
    first.additional_information = {"meta": {"codec_streaming": True}, "nvc_logical_prompt_len": 1}
    first_frame = torch.full((1, 31), 101, dtype=torch.long)

    talker2code2wav_async_chunk(
        adapter,
        {"codes": {"audio": first_frame}, "meta": {"codec_streaming": True}},
        first,
        is_finished=False,
    )

    adapter.requests_num_chunks_sent[external_req_id] = 1
    adapter.finish_requests([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first})

    assert external_req_id not in adapter.request_payload
    assert external_req_id not in adapter.requests_num_chunks_sent

    replacement = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
    replacement.resumable = True
    replacement.model_intermediate_buffer = None
    replacement.additional_information = first.additional_information
    replacement_payload = talker2code2wav_async_chunk(
        adapter,
        {"codes": {"audio": torch.full((1, 31), 202, dtype=torch.long)}, "meta": {"codec_streaming": True}},
        replacement,
        is_finished=False,
    )

    assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long))


def test_save_async_boundary_holds_generation_without_request_counter(build_adapter):
    """A boundary must not advance past a generation nothing can catch up to.

    When the producing stage maintains ``_omni_segment_generation`` on the
    request, a boundary may start the next segment's watermark at
    ``generation + 1`` because the following frame will carry it. On paths
    where the request has no such counter, that speculative advance makes every
    subsequent frame look stale and starves the stream permanently
    (vllm-project/vllm-omni#6816).
    """
    adapter, _ = build_adapter(stage_id=1)

    request = _req("req-nogen", RequestStatus.WAITING, external_req_id="ext-nogen")
    request.resumable = True
    request._omni_segment_generation = None

    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
    assert adapter._segment_generation["ext-nogen"] == 0

    follow_up = _req("req-nogen-next", RequestStatus.WAITING, external_req_id="ext-nogen")
    follow_up.resumable = True
    follow_up.num_computed_tokens = 3
    follow_up._omni_segment_generation = None
    queued_before = len(adapter._pending_save_reqs)

    adapter.save_async(multimodal_output=None, request=follow_up, is_segment_finished=False)

    assert len(adapter._pending_save_reqs) == queued_before + 1
    assert adapter._pending_save_reqs[-1]["request"] is follow_up