vllm-omni/tests/engine/test_output_processor.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""

from dataclasses import dataclass
from types import SimpleNamespace
from unittest.mock import MagicMock

import numpy as np
import pytest
import torch
from vllm.outputs import PoolingRequestOutput
from vllm.sampling_params import RequestOutputKind
from vllm.v1.engine import EngineCoreEvent, EngineCoreEventType, FinishReason
from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
from vllm.v1.metrics.stats import IterationStats, PrefillStats

from vllm_omni.engine import OmniEngineCoreOutput
from vllm_omni.outputs import output_processor
from vllm_omni.outputs.output_modality import OutputModality, OutputModalityNames
from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState

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

# Audio is explicitly listed as a drainable modality
AUDIO = OutputModalityNames.AUDIO

# Latent is explicitly not drainable, but the choice doesn't matter here as
# long as isn't listed as drainable. I.e., could also be arbitrary keys for
# the purposes of these tests
LATENT = OutputModalityNames.LATENT

# NOTE: detokenizer and logprobs aren't really used here, but we mock them since
# some of the utils called in vLLM superclass assert require them to be None.
_DETOK = MagicMock(
    output_token_ids=[0],
    get_next_output_text=MagicMock(return_value=""),
    num_output_tokens=MagicMock(return_value=1),
)
_LOGPROBS = MagicMock(logprobs=None, cumulative_logprob=None, prompt_logprobs=None)


@dataclass
class _StreamingUpdate:
    final: bool
    prompt: object | None
    prompt_token_ids: list[int]
    arrival_time: float


@dataclass
class _FinishedRequestMetric:
    request_id: str
    mean_time_per_output_token: float


_DEFAULT_STATE_KWARGS = dict(
    request_id="r",
    external_req_id="r",
    parent_req=None,
    request_index=0,
    lora_request=None,
    prompt=None,
    prompt_token_ids=[0],
    prompt_embeds=None,
    logprobs_processor=_LOGPROBS,
    detokenizer=_DETOK,
    max_tokens_param=None,
    arrival_time=0.0,
    queue=None,
    log_stats=False,
    stream_interval=1,
)


def _make_state(output_kind: RequestOutputKind):
    return OmniRequestState(**_DEFAULT_STATE_KWARGS, output_kind=output_kind)


def test_output_modality_name_uses_string_value():
    assert str(OutputModalityNames.AUDIO) == "audio"


def test_init_empty_dict():
    """Ensure mm_accumulated is initially empty."""
    assert _make_state(RequestOutputKind.CUMULATIVE).mm_accumulated == {}
    assert _make_state(RequestOutputKind.DELTA).mm_accumulated == {}


def test_streaming_update_resets_native_text_metrics_for_next_segment():
    state = _make_state(RequestOutputKind.DELTA)
    state.native_text_stats.num_generation_tokens = 12
    state.native_text_stats.first_token_ts = 10.0
    state.native_text_stats.last_token_ts = 10.2

    state.apply_streaming_update(
        SimpleNamespace(
            final=False,
            prompt=None,
            prompt_token_ids=[],
            arrival_time=20.0,
        )
    )

    assert state.native_text_stats.arrival_time == 20.0
    assert state.native_text_stats.num_generation_tokens == 0
    assert state.native_text_stats.first_token_ts == 0.0
    assert state.native_text_stats.last_token_ts == 0.0


def test_native_text_tpot_snapshot_matches_vllm_finished_metric_definition():
    calculate = getattr(output_processor, "_mean_time_per_output_token_ms", None)
    assert calculate is not None
    stats = SimpleNamespace(
        num_generation_tokens=4,
        first_token_ts=10.0,
        last_token_ts=10.045,
    )

    assert calculate(stats) == pytest.approx(15.0)


def test_native_text_metrics_include_segment_generation_token_count(monkeypatch):
    monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
    processor = object.__new__(MultimodalOutputProcessor)
    processor._native_text_metrics_by_request = {}
    processor.lora_states = {}
    state = _make_state(RequestOutputKind.DELTA)
    state.is_prefilling = False
    iteration_stats = MagicMock()

    def update_native_stats(_output, _timestamp, _was_prefilling, native_stats, *_args):
        native_stats.num_generation_tokens = 27
        native_stats.first_token_ts = 10.0
        native_stats.last_token_ts = 10.26

    iteration_stats.update_from_output.side_effect = update_native_stats

    processor._update_stats_from_output(
        state,
        SimpleNamespace(),
        10.26,
        iteration_stats,
    )

    assert processor.pop_native_text_metrics("r")["num_generation_tokens"] == 27


def test_native_text_metrics_exclude_cross_segment_gap(monkeypatch):
    monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
    processor = object.__new__(MultimodalOutputProcessor)
    processor._native_text_metrics_by_request = {}
    processor.lora_states = {}
    state = _make_state(RequestOutputKind.DELTA)
    iteration_stats = MagicMock()

    def update_segment(_output, timestamp, _was_prefilling, native_stats, *_args):
        native_stats.num_generation_tokens = 2 if timestamp == 10.03 else 1
        native_stats.first_token_ts = 10.02 if timestamp == 10.03 else timestamp
        native_stats.last_token_ts = timestamp
        native_stats.first_token_latency = 0.01

    iteration_stats.update_from_output.side_effect = update_segment
    processor._update_stats_from_output(
        state,
        MagicMock(),
        10.0,
        iteration_stats,
    )
    first_segment = processor.pop_native_text_metrics("r")
    assert first_segment["vllm_tpot_ms"] == 0.0
    assert first_segment["vllm_itls_ms"] == []

    state.apply_streaming_update(
        _StreamingUpdate(
            final=False,
            prompt=None,
            prompt_token_ids=[],
            arrival_time=10.01,
        )
    )
    state.is_prefilling = True
    processor._update_stats_from_output(
        state,
        MagicMock(),
        10.02,
        iteration_stats,
    )
    state.is_prefilling = False
    processor._update_stats_from_output(
        state,
        MagicMock(),
        10.03,
        iteration_stats,
    )
    second_segment = processor.pop_native_text_metrics("r")

    assert second_segment["num_generation_tokens"] == 2
    assert second_segment["vllm_itls_ms"] == pytest.approx([10.0])
    assert second_segment["vllm_tpot_ms"] == pytest.approx(10.0)


def test_native_text_tpot_weights_multi_token_engine_output(monkeypatch):
    monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
    processor = object.__new__(MultimodalOutputProcessor)
    processor._native_text_metrics_by_request = {}
    processor.lora_states = {}
    state = _make_state(RequestOutputKind.DELTA)
    iteration_stats = MagicMock()

    def update_segment(_output, timestamp, _was_prefilling, native_stats, *_args):
        native_stats.num_generation_tokens = 1 if timestamp == 10.0 else 4
        native_stats.first_token_ts = 10.0
        native_stats.last_token_ts = timestamp
        native_stats.first_token_latency = 0.01

    iteration_stats.update_from_output.side_effect = update_segment
    state.is_prefilling = True
    processor._update_stats_from_output(state, MagicMock(), 10.0, iteration_stats)
    state.is_prefilling = False
    processor._update_stats_from_output(state, MagicMock(), 10.03, iteration_stats)
    record = processor.pop_native_text_metrics("r")

    assert record["num_generation_tokens"] == 4
    assert record["vllm_itls_ms"] == pytest.approx([30.0])
    assert record["vllm_itl_ms"] == pytest.approx(30.0)
    assert record["vllm_tpot_ms"] == pytest.approx(10.0)
    assert "_tpot_elapsed_ms" not in record
    assert "_tpot_intervals" not in record


@pytest.mark.parametrize(
    ("finished_tpot_s", "expected_tpot_ms"),
    [(0.0, 19.0), (0.023, 23.0)],
)
def test_native_text_tpot_only_accepts_positive_finished_metric(
    monkeypatch,
    finished_tpot_s,
    expected_tpot_ms,
):
    monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_finished", lambda *args, **kwargs: None)
    processor = object.__new__(MultimodalOutputProcessor)
    processor._native_text_metrics_by_request = {"r": {"vllm_tpot_ms": 19.0}}
    state = _make_state(RequestOutputKind.DELTA)
    iteration_stats = MagicMock()
    iteration_stats.finished_requests = [
        _FinishedRequestMetric(
            request_id="r",
            mean_time_per_output_token=finished_tpot_s,
        )
    ]

    processor._update_stats_from_finished(
        state,
        FinishReason.STOP,
        iteration_stats,
    )

    assert processor.pop_native_text_metrics("r")["vllm_tpot_ms"] == expected_tpot_ms


def test_suppressed_native_tokens_use_authoritative_count_even_without_stats(monkeypatch):
    monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
    processor = object.__new__(MultimodalOutputProcessor)
    processor._native_text_metrics_by_request = {}
    state = _make_state(RequestOutputKind.FINAL_ONLY)
    processor._update_stats_from_output(
        state, SimpleNamespace(new_token_ids=[2150], num_generation_tokens=2048), None, None
    )
    assert processor.pop_native_text_metrics("r")["num_generation_tokens"] == 2048


def test_delta_drains_output_modality_per_step():
    """DELTA drains the mm_type key (output modality) but preserves hidden-state keys."""
    s = _make_state(RequestOutputKind.DELTA)
    audio1, audio2, hs1, hs2 = [torch.ones(num_elem) for num_elem in range(1, 5)]

    # Add audio and hidden state tensors
    s.add_multimodal_tensor(audio1, mm_type=AUDIO)  # should be drained
    s.add_multimodal_tensor(hs1, mm_type=LATENT)  # shouldn't be drained

    out1 = s._new_completion_output([1], None, None)
    out1_audio = out1.multimodal_output[AUDIO]
    out1_hidden = out1.multimodal_output[LATENT]
    assert isinstance(out1_audio, torch.Tensor)
    assert torch.equal(out1.multimodal_output[AUDIO], audio1)
    assert isinstance(out1_hidden, torch.Tensor)
    assert torch.equal(out1.multimodal_output[LATENT], hs1)

    # After emission, hidden states should remain, but audio is drained
    assert set(s.mm_accumulated.keys()) == {LATENT}

    s.add_multimodal_tensor(audio2, AUDIO)
    s.add_multimodal_tensor(hs2, mm_type=LATENT)
    out2 = s._new_completion_output([2], None, None)
    out2_audio = out2.multimodal_output[AUDIO]
    out2_hidden = out2.multimodal_output[LATENT]
    assert isinstance(out2_audio, torch.Tensor)
    assert torch.equal(out2_audio, audio2)
    # Since hidden isn't drained, it's grown to a list
    assert isinstance(out2_hidden, list) and len(out2_hidden) == 2
    assert torch.equal(hs1, out2_hidden[0])
    assert torch.equal(hs2, out2_hidden[1])


def test_cumulative_emits_consolidated_audio_each_step():
    """Ensure cumulative accumulates and consolidates modality keys every step."""
    s = _make_state(RequestOutputKind.CUMULATIVE)
    # NOTE: audio is usually emitted as (1, size) chunks; we need to be sure
    # to not change the tensor dimension when we consolidate
    audio1 = torch.ones(1, 500)
    s.add_multimodal_tensor(audio1, mm_type=AUDIO)
    req_out = s.make_request_output([1], None, None, None)
    assert req_out is not None
    cons_audio = req_out.outputs[0].multimodal_output[AUDIO]
    # Single chunk keeps original shape [1, 500]
    assert isinstance(cons_audio, torch.Tensor) and cons_audio.shape == audio1.shape

    audio2 = torch.ones(1, 300)
    s.add_multimodal_tensor(audio2, mm_type=AUDIO)
    req_out = s.make_request_output([2], None, None, None)
    assert req_out is not None
    cons_audio = req_out.outputs[0].multimodal_output[AUDIO]
    # After consolidation, audio chunks are concatenated on last axis,
    # preserving the [1, N] channel dimension
    total_audio_len = audio1.shape[-1] + audio2.shape[-1]
    assert isinstance(audio2, torch.Tensor) and cons_audio.shape == (1, total_audio_len)

    assert "audio" in s.mm_accumulated


def test_finish_consolidates_hidden_states():
    """Ensure consolidation merges hidden-state tensor lists on finish."""
    s = _make_state(RequestOutputKind.CUMULATIVE)
    s.add_multimodal_tensor(torch.ones(5, 4), mm_type=LATENT)
    s.add_multimodal_tensor(torch.ones(3, 4), mm_type=LATENT)

    result = s.make_request_output([1], None, FinishReason.STOP, None)
    assert result is not None and not isinstance(result, PoolingRequestOutput)

    hs = result.outputs[0].multimodal_output[LATENT]
    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 8


def test_finish_consolidation_for_hs_delta():
    """Ensure finish doesn't drop the accumulated hidden states."""
    s = _make_state(RequestOutputKind.DELTA)
    # hidden state accumulation (nothing drained)
    s.add_multimodal_tensor({"foo": torch.ones(5, 4)}, mm_type=LATENT)
    result = s.make_request_output([0], None, FinishReason.STOP, None)
    assert result is not None and not isinstance(result, PoolingRequestOutput)
    hs = result.outputs[0].multimodal_output["foo"]
    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 5

    # Since we don't drain the hidden states, if we add 3 elements, we should get 8
    s.add_multimodal_tensor({"foo": torch.ones(3, 4)}, mm_type=LATENT)
    result = s.make_request_output([0], None, FinishReason.STOP, None)
    assert result is not None and not isinstance(result, PoolingRequestOutput)
    hs = result.outputs[0].multimodal_output["foo"]
    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 8
    assert "foo" in s.mm_accumulated


def test_finish_consolidation_drains_mm_delta():
    """Ensure making the request output drains modality deltas (e.g., audio)."""
    s = _make_state(RequestOutputKind.DELTA)
    # multimodal data accumulation (drained)
    s.add_multimodal_tensor({AUDIO: torch.ones(5, 4)}, mm_type=AUDIO)
    result = s.make_request_output([0], None, FinishReason.STOP, None)
    assert result is not None and not isinstance(result, PoolingRequestOutput)
    hs = result.outputs[0].multimodal_output[AUDIO]
    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 5

    # Since we did drain the hidden states, we no longer get the 5 back
    s.add_multimodal_tensor({AUDIO: torch.ones(3, 4)}, mm_type=AUDIO)
    result = s.make_request_output([0], None, FinishReason.STOP, None)
    assert result is not None and not isinstance(result, PoolingRequestOutput)
    hs = result.outputs[0].multimodal_output[AUDIO]
    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 3
    assert AUDIO not in s.mm_accumulated  # drained


def test_delta_drains_audio_chunk_metadata_per_step():
    s = _make_state(RequestOutputKind.DELTA)
    segment_text = "你好"
    segment_utf8 = torch.tensor(list(segment_text.encode("utf-8")), dtype=torch.uint8)

    s.add_multimodal_tensor(
        {
            "model_outputs": torch.ones(5),
            "meta.duplex_epoch": torch.tensor([7], dtype=torch.int32),
            "meta.duplex_turn_id": torch.tensor([1], dtype=torch.int32),
            "meta.llm_output_text_utf8": segment_utf8,
            "meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
            "meta.tts_is_last_chunk": torch.tensor([0], dtype=torch.int32),
        },
        mm_type=AUDIO,
    )
    result1 = s.make_request_output([1], None, FinishReason.STOP, None)
    assert result1 is not None and not isinstance(result1, PoolingRequestOutput)
    meta1 = result1.outputs[0].multimodal_output["meta"]
    assert bytes(meta1["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text

    s.add_multimodal_tensor(
        {
            "model_outputs": torch.ones(3),
            "meta.duplex_epoch": torch.tensor([8], dtype=torch.int32),
            "meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32),
            "meta.llm_output_text_utf8": segment_utf8,
            "meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
            "meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32),
        },
        mm_type=AUDIO,
    )
    result2 = s.make_request_output([2], None, FinishReason.STOP, None)

    assert result2 is not None and not isinstance(result2, PoolingRequestOutput)
    meta2 = result2.outputs[0].multimodal_output["meta"]
    assert bytes(meta2["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text
    assert meta2["audio_text_total_chars"].tolist() == [len(segment_text)]
    assert meta2["duplex_epoch"].tolist() == [8]
    assert meta2["duplex_turn_id"].tolist() == [2]


def test_cumulative_audio_replaces_chunk_metadata_per_step():
    s = _make_state(RequestOutputKind.CUMULATIVE)
    segment_text = "你好"
    segment_utf8 = torch.tensor(list(segment_text.encode("utf-8")), dtype=torch.uint8)

    s.add_multimodal_tensor(
        {
            "model_outputs": torch.ones(5),
            "meta.duplex_epoch": torch.tensor([7], dtype=torch.int32),
            "meta.duplex_turn_id": torch.tensor([1], dtype=torch.int32),
            "meta.llm_output_text_utf8": segment_utf8,
            "meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
            "meta.tts_is_last_chunk": torch.tensor([0], dtype=torch.int32),
        },
        mm_type=AUDIO,
    )
    result1 = s.make_request_output([1], None, None, None)
    assert result1 is not None and not isinstance(result1, PoolingRequestOutput)
    assert result1.outputs[0].multimodal_output["audio"].shape[0] == 5

    s.add_multimodal_tensor(
        {
            "model_outputs": torch.ones(3),
            "meta.duplex_epoch": torch.tensor([8], dtype=torch.int32),
            "meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32),
            "meta.llm_output_text_utf8": segment_utf8,
            "meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
            "meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32),
        },
        mm_type=AUDIO,
    )
    result2 = s.make_request_output([2], None, None, None)

    assert result2 is not None and not isinstance(result2, PoolingRequestOutput)
    output = result2.outputs[0].multimodal_output
    assert output["audio"].shape[0] == 8
    meta = output["meta"]
    assert bytes(meta["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text
    assert meta["audio_text_total_chars"].tolist() == [len(segment_text)]
    assert meta["duplex_epoch"].tolist() == [8]
    assert meta["duplex_turn_id"].tolist() == [2]
    assert meta["tts_is_last_chunk"].tolist() == [1]


def test_delta_audio_non_final_tts_chunk_overrides_spurious_finish():
    s = _make_state(RequestOutputKind.DELTA)
    s.add_multimodal_tensor(
        {
            "model_outputs": torch.ones(5),
            "meta.tts_is_last_chunk": torch.tensor([0]),
        },
        mm_type=AUDIO,
    )

    result = s.make_request_output([1], None, FinishReason.STOP, None)

    assert result is not None and not isinstance(result, PoolingRequestOutput)
    assert result.finished is False
    assert result.outputs[0].finish_reason is None
    assert result.outputs[0].multimodal_output["meta"]["tts_is_last_chunk"].tolist() == [0]
    assert torch.equal(result.outputs[0].multimodal_output[AUDIO], torch.ones(5))


def test_delta_audio_final_tts_chunk_keeps_finish():
    s = _make_state(RequestOutputKind.DELTA)
    s.add_multimodal_tensor(
        {
            "model_outputs": torch.ones(5),
            "meta.tts_is_last_chunk": torch.tensor([1]),
        },
        mm_type=AUDIO,
    )

    result = s.make_request_output([151645], None, FinishReason.STOP, None)

    assert result is not None and not isinstance(result, PoolingRequestOutput)
    assert result.finished is True
    assert result.outputs[0].finish_reason == "stop"


def test_delta_audio_non_final_tts_chunk_accepts_nested_meta():
    s = _make_state(RequestOutputKind.DELTA)
    s.add_multimodal_tensor(
        {
            AUDIO: torch.ones(5),
            "meta": {"tts_is_last_chunk": torch.tensor([0])},
        },
        mm_type=AUDIO,
    )

    result = s.make_request_output([1], None, FinishReason.STOP, None)

    assert result is not None and not isinstance(result, PoolingRequestOutput)
    assert result.finished is False
    assert result.outputs[0].multimodal_output["meta"]["tts_is_last_chunk"].tolist() == [0]


@pytest.mark.parametrize("mm_type", [AUDIO, "hidden"])
def test_final_only_consolidates_drainable_keys(mm_type):
    """FINAL_ONLY never drains per-step, so modality keys and hidden state
    keys both accumulate and are consolidated on finish."""
    s = _make_state(RequestOutputKind.FINAL_ONLY)

    # NOTE: Currently there is brittlness in the tensor stacking, so we just
    # test a 1D tensor here. The intention is just to ensure audio /hidden
    # behave the same.
    s.add_multimodal_tensor(torch.ones(500), mm_type=mm_type)
    # Non-finish step returns None without calling _new_completion_output
    assert s.make_request_output([1], None, None, None) is None
    assert mm_type in s.mm_accumulated

    s.add_multimodal_tensor(torch.ones(300), mm_type=mm_type)
    result = s.make_request_output([2], None, FinishReason.STOP, None)
    assert result is not None and not isinstance(result, PoolingRequestOutput)

    audio = result.outputs[0].multimodal_output[mm_type]
    assert isinstance(audio, torch.Tensor)
    assert audio.shape == (800,)


def test_cumulative_token_ids_always_set():
    """cumulative_token_ids is set for all output kinds."""
    for kind in (RequestOutputKind.DELTA, RequestOutputKind.CUMULATIVE, RequestOutputKind.FINAL_ONLY):
        s = _make_state(kind)
        out = s._new_completion_output([42], None, None)
        assert hasattr(out, "cumulative_token_ids")
        # The mock detokenizer has output_token_ids=[0]
        assert list(out.cumulative_token_ids) == [0]


def test_cumulative_token_ids_is_a_copy():
    """cumulative_token_ids must be a snapshot, not a live reference."""
    s = _make_state(RequestOutputKind.DELTA)
    out = s._new_completion_output([42], None, None)
    assert out.cumulative_token_ids is not _DETOK.output_token_ids


# ---------------------------------------------------------------------------
# No-detokenizer regression tests (Concern 2)
# ---------------------------------------------------------------------------

_NO_DETOK_STATE_KWARGS = {
    **_DEFAULT_STATE_KWARGS,
    "detokenizer": None,
    "logprobs_processor": None,
}


def _make_no_detok_state(output_kind: RequestOutputKind):
    return OmniRequestState(**_NO_DETOK_STATE_KWARGS, output_kind=output_kind)


def test_no_detokenizer_completion_output():
    """_new_completion_output works when detokenizer is None (generation stages)."""
    s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
    out = s._new_completion_output([1, 2], None, None)
    # Should still produce output with empty text and the token_ids as-is
    assert out.text == ""
    assert list(out.token_ids) == [1, 2]
    assert hasattr(out, "cumulative_token_ids")
    assert list(out.cumulative_token_ids) == [1, 2]
    # Multimodal data should be attached
    assert AUDIO in out.multimodal_output


def test_no_detokenizer_make_request_output():
    """make_request_output works without detokenizer when multimodal data is present."""
    s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
    result = s.make_request_output([], None, FinishReason.STOP, None)
    assert result is not None
    assert not isinstance(result, PoolingRequestOutput)
    assert AUDIO in result.outputs[0].multimodal_output


def test_no_detokenizer_process_outputs_returns_nonterminal_audio_chunk(monkeypatch):
    """Generation-stage MM-only chunks must reach the orchestrator.

    StagePool registers generation requests without a per-request collector,
    so the processor return value is the only path from Stage2 to
    Orchestrator._handle_processed_outputs().
    """
    state = _make_no_detok_state(RequestOutputKind.DELTA)
    processor = object.__new__(MultimodalOutputProcessor)
    processor.output_modality = OutputModality.AUDIO
    processor.request_states = {"r": state}
    upstream_result = SimpleNamespace(request_outputs=[], reqs_to_abort=[])
    monkeypatch.setattr(
        VLLMOutputProcessor,
        "process_outputs",
        lambda *_args, **_kwargs: upstream_result,
    )
    engine_output = SimpleNamespace(
        request_id="r",
        multimodal_output={
            "model_outputs": torch.arange(8, dtype=torch.float32),
            "sr": torch.tensor(24000, dtype=torch.int32),
        },
        output_type="audio",
        pooling_output=None,
        new_token_ids=[],
        finish_reason=None,
        stop_reason=None,
        kv_transfer_params=None,
        ec_transfer_params=None,
        routed_experts=None,
        num_cached_tokens=0,
    )

    result = processor.process_outputs([engine_output])

    assert len(result.request_outputs) == 1
    assert result.request_outputs[0].request_id == "r"
    assert AUDIO in result.request_outputs[0].outputs[0].multimodal_output


def _make_mm_only_output_processor(monkeypatch):
    processor = object.__new__(MultimodalOutputProcessor)
    processor.output_modality = OutputModality.AUDIO
    processor.request_states = {"r": _make_no_detok_state(RequestOutputKind.DELTA)}
    processor.tracing_enabled = False
    monkeypatch.setattr(
        VLLMOutputProcessor,
        "process_outputs",
        lambda *_args, **_kwargs: SimpleNamespace(
            request_outputs=[],
            reqs_to_abort=[],
        ),
    )

    def finish_request(req_state):
        processor.request_states.pop(req_state.request_id, None)

    processor._finish_request = finish_request
    return processor


def _audio_engine_output(*, is_segment_finished: bool, is_last_chunk: bool):
    return SimpleNamespace(
        request_id="r",
        multimodal_output={
            "model_outputs": torch.arange(8, dtype=torch.float32),
            "meta.tts_is_last_chunk": torch.tensor([int(is_last_chunk)]),
        },
        output_type="audio",
        pooling_output=None,
        new_token_ids=[],
        finish_reason=FinishReason.STOP,
        stop_reason=None,
        kv_transfer_params=None,
        ec_transfer_params=None,
        routed_experts=None,
        num_cached_tokens=0,
        is_segment_finished=is_segment_finished,
    )


def test_mm_only_segment_finish_retains_request_state_for_next_audio_chunk(monkeypatch):
    processor = _make_mm_only_output_processor(monkeypatch)

    first = processor.process_outputs([_audio_engine_output(is_segment_finished=True, is_last_chunk=False)])

    assert "r" in processor.request_states
    assert len(first.request_outputs) == 1
    assert first.request_outputs[0].finished is False

    second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])

    assert len(second.request_outputs) == 1
    assert second.request_outputs[0].finished is True
    assert "r" not in processor.request_states


def test_mm_only_nonfinal_audio_retains_state_without_segment_marker(monkeypatch):
    processor = _make_mm_only_output_processor(monkeypatch)

    first = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=False)])

    assert "r" in processor.request_states
    assert len(first.request_outputs) == 1
    assert first.request_outputs[0].finished is False

    second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])

    assert len(second.request_outputs) == 1
    assert second.request_outputs[0].finished is True
    assert "r" not in processor.request_states


def test_mm_only_terminal_finish_removes_request_state(monkeypatch):
    processor = _make_mm_only_output_processor(monkeypatch)

    result = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])

    assert len(result.request_outputs) == 1
    assert result.request_outputs[0].finished is True
    assert "r" not in processor.request_states


def test_no_detokenizer_make_request_output_with_routed_experts():
    """make_request_output accepts the routed_experts arg that the multimodal
    output channel (_process_mm_only_outputs) passes, and attaches it to the
    completion output.

    routed_experts is omni-specific keyword-only: the 6th positional now
    matches upstream's ec_transfer_params so that super().process_outputs()
    calls do not misroute it.
    """
    s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
    routed_experts = np.zeros((2, 3), dtype=np.int32)
    # Mirror the exact call shape of _process_mm_only_outputs.
    result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts=routed_experts)
    assert result is not None
    assert not isinstance(result, PoolingRequestOutput)
    assert result.outputs[0].routed_experts is routed_experts


def test_no_detokenizer_stream_interval_skipped():
    """stream_interval > 1 does not crash when detokenizer is None."""
    kwargs = {**_NO_DETOK_STATE_KWARGS, "stream_interval": 5}
    s = OmniRequestState(**kwargs, output_kind=RequestOutputKind.DELTA)
    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
    # Should not assert; stream_interval logic is simply skipped
    result = s.make_request_output([], None, FinishReason.STOP, None)
    assert result is not None


def test_no_detokenizer_final_only():
    """FINAL_ONLY with no detokenizer returns None until finish."""
    s = _make_no_detok_state(RequestOutputKind.FINAL_ONLY)
    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
    # Non-finish step should return None
    assert s.make_request_output([], None, None, None) is None
    # Finish step should return output
    result = s.make_request_output([], None, FinishReason.STOP, None)
    assert result is not None
    assert AUDIO in result.outputs[0].multimodal_output


def test_mm_only_outputs_update_iteration_stats():
    processor = MultimodalOutputProcessor(
        tokenizer=None,
        log_stats=True,
        engine_core_output_type=AUDIO,
    )
    state = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
    processor.request_states[state.request_id] = state
    processor.external_req_ids[state.external_req_id].append(state.request_id)

    prefill_stats = PrefillStats()
    prefill_stats.set(
        num_prompt_tokens=4,
        num_local_cached_tokens=0,
        num_external_cached_tokens=0,
    )
    output = OmniEngineCoreOutput(
        request_id=state.request_id,
        new_token_ids=[10, 11],
        finish_reason=FinishReason.STOP,
        events=[
            EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0),
            EngineCoreEvent(EngineCoreEventType.SCHEDULED, 2.0),
        ],
        prefill_stats=prefill_stats,
        multimodal_output={AUDIO: torch.ones(1, 8)},
    )
    iteration_stats = IterationStats()

    processor.process_outputs(
        [output],
        engine_core_timestamp=3.0,
        iteration_stats=iteration_stats,
    )

    assert iteration_stats.num_prompt_tokens == 4
    assert iteration_stats.num_generation_tokens == 2
    assert len(iteration_stats.time_to_first_tokens_iter) == 1
    assert len(iteration_stats.finished_requests) == 1
    finished = iteration_stats.finished_requests[0]
    assert finished.finish_reason == FinishReason.STOP
    assert finished.num_prompt_tokens == state.prompt_len
    assert finished.num_generation_tokens == 2


def test_ec_transfer_params_survive_both_construction_paths():
    """v0.28 contract: encoder-cache transfer metadata must reach
    RequestOutput.ec_transfer_params through BOTH omni construction paths —
    the upstream-delegating one (logprobs processor present) and the
    no-detokenizer generation-stage one (direct RequestOutput build).
    Regression: the override accepted the argument but dropped it."""
    ec = {"remote_handle": "enc-cache-1"}

    state = _make_state(RequestOutputKind.CUMULATIVE)
    out = state.make_request_output([1], None, None, None, {"kv": 1}, ec)
    assert out is not None and out.ec_transfer_params == ec

    kwargs = dict(_DEFAULT_STATE_KWARGS)
    kwargs.update(logprobs_processor=None, detokenizer=None)
    gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
    completion = gen_state._new_completion_output([1], None, None)
    direct = gen_state._new_request_output("r", [completion], False, {"kv": 1}, ec)
    assert direct.ec_transfer_params == ec


def test_num_cache_creation_tokens_reaches_direct_request_output():
    """v0.28 contract: prefix-cache creation usage must reach
    RequestOutput.num_cache_creation_tokens through the no-detokenizer
    direct build, which previously passed only num_cached_tokens."""
    kwargs = dict(_DEFAULT_STATE_KWARGS)
    kwargs.update(logprobs_processor=None, detokenizer=None)
    gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
    gen_state.num_cached_tokens = 3
    gen_state.num_cache_creation_tokens = 2

    completion = gen_state._new_completion_output([1], None, None)
    direct = gen_state._new_request_output("r", [completion], False, None, None)

    assert direct.num_cached_tokens == 3
    assert direct.num_cache_creation_tokens == 2


def _abort_processor_with_parent(child_ids: list[str]):
    processor = object.__new__(MultimodalOutputProcessor)
    processor.request_states = {}
    processor.external_req_ids = {}
    processor.parent_requests = {}
    processor._native_text_metrics_by_request = {}
    processor.lora_states = SimpleNamespace(request_finished=lambda *_args, **_kwargs: None)
    parent = SimpleNamespace(request_id="parent", child_requests=set(child_ids))
    processor.parent_requests["parent"] = parent
    for child_id in child_ids:
        req_state = SimpleNamespace(
            parent_req=parent,
            lora_name=None,
            output_kind=RequestOutputKind.CUMULATIVE,
            detokenizer=SimpleNamespace(output_token_ids=[7, 8]),
            queue=None,
            external_req_id="parent",
            make_request_output=MagicMock(return_value=SimpleNamespace(request_id=child_id)),
        )
        processor.request_states[child_id] = req_state
    return processor, parent


def test_abort_last_parallel_child_drops_parent_request():
    processor, parent = _abort_processor_with_parent(["0_parent", "1_parent"])

    aborted, _outputs = processor.abort_requests_collecting_outputs(["0_parent"], internal=True)
    assert aborted == ["0_parent"]
    assert "parent" in processor.parent_requests
    assert parent.child_requests == {"1_parent"}

    aborted, _outputs = processor.abort_requests_collecting_outputs(["1_parent"], internal=True)
    assert aborted == ["1_parent"]
    assert "parent" not in processor.parent_requests
    assert not parent.child_requests


def test_abort_snapshot_leaves_state_until_commit():
    processor, parent = _abort_processor_with_parent(["0_parent"])
    processor.external_req_ids["parent"] = ["0_parent"]

    aborted, outputs = processor.abort_requests_collecting_outputs(
        ["parent"],
        internal=False,
        commit_state=False,
    )
    assert aborted == ["0_parent"]
    assert outputs
    assert "0_parent" in processor.request_states
    assert processor.external_req_ids["parent"] == ["0_parent"]
    assert "parent" in processor.parent_requests
    assert parent.child_requests == {"0_parent"}

    processor.commit_aborted_request_state(["parent"], internal=False)
    assert "0_parent" not in processor.request_states
    assert "parent" not in processor.parent_requests
    assert "parent" not in processor.external_req_ids


def test_cumulative_audio_terminal_output_supplies_full_aligner_waveform():
    from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner

    state = _make_state(RequestOutputKind.CUMULATIVE)
    for index, count in enumerate((800, 1200, 400)):
        state.add_multimodal_tensor(torch.full((1, count), float(index)), mm_type=AUDIO)
        output = state.make_request_output([index], None, FinishReason.STOP if index == 2 else None, None)
        assert output is not None
        if index < 2:
            assert code2wav2aligner([output], {}) == []
    prompt = {"additional_information": {"text": ["Hello world"], "language": ["English"]}}
    aligned = code2wav2aligner([output], prompt)[0]
    waveform, sr = aligned["multi_modal_data"]["audio"]
    assert sr == 24000
    assert waveform.shape == (2400,)
    assert (waveform[:800] == 0).all()
    assert (waveform[800:2000] == 1).all()
    assert (waveform[2000:] == 2).all()
    assert aligned["additional_information"]["aligner_audio_duration_ms"] == [100.0]