vllm-omni/tests/entrypoints/openai_api/test_codec_stream_protocol.py

"""Codec-token streaming: wire format, row alignment, and handler behaviour (CPU)."""

import numpy as np
import pytest
import torch
from fastapi import FastAPI, WebSocket

from vllm_omni.entrypoints.openai import codec_stream as cs
from vllm_omni.entrypoints.openai.protocol.audio import StreamingSpeechSessionConfig

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


def test_pack_unpack_roundtrip_and_flags():
    codes = np.random.randint(0, 2048, size=(25, 16)).astype(np.int64)
    buf = cs.pack_codec_frame(7, torch.from_numpy(codes), eos=False, segment_end=True)
    assert len(buf) == cs.HEADER_SIZE + 25 * 16 * 2
    seq, n, flags, arr = cs.unpack_codec_frame(buf, 16)
    assert seq == 7 and n == 25 and flags == cs.FLAG_SEGMENT_END
    assert arr.dtype == np.int16 and np.array_equal(arr.astype(np.int64), codes)
    # EOS with zero frames
    buf = cs.pack_codec_frame(8, None, eos=True, codebooks=16)
    seq, n, flags, arr = cs.unpack_codec_frame(buf, 16)
    assert (seq, n, flags) == (8, 0, cs.FLAG_EOS) and arr.shape == (0, 16)
    buf = cs.pack_codec_frame(2**32 + 5, np.zeros((1, 16), dtype=np.int64), reset=True)
    assert cs.unpack_codec_frame(buf, 16)[0] == 5
    assert cs.unpack_codec_frame(buf, 16)[2] == cs.FLAG_RESET


def test_pack_rejects_bad_shapes_and_ranges():
    with pytest.raises(ValueError):
        cs.pack_codec_frame(0, np.zeros((2, 2, 2), dtype=np.int64))
    with pytest.raises(ValueError):
        cs.pack_codec_frame(0, np.full((1, 16), 40000, dtype=np.int64))
    with pytest.raises(ValueError):
        cs.unpack_codec_frame(b"\x00\x01", 16)
    with pytest.raises(ValueError):
        cs.unpack_codec_frame(cs.pack_codec_frame(0, np.zeros((1, 16), dtype=np.int64))[:-2], 16)


def test_align_codec_rows_drops_prefill_placeholder_and_eos():
    # 1 prefill placeholder row (zeros, emitted before any token) + 3 real frames + EOS row
    codes = torch.tensor([[0] * 16, [5] * 16, [6] * 16, [7] * 16, [0] * 16])
    token_ids = [5, 6, 7, 2150]  # codebook-0 ids then EOS
    rows = cs.align_codec_rows(codes, token_ids, 2048)
    assert rows.shape == (3, 16) and rows[:, 0].tolist() == [5, 6, 7]
    # trailing -1 placeholders in token ids are ignored
    rows = cs.align_codec_rows(codes, token_ids + [-1, -1], 2048)
    assert rows.shape == (3, 16)
    # out-of-range rows are dropped
    bad = torch.tensor([[5] * 16, [3000] * 16])
    assert cs.align_codec_rows(bad, [5, 6], 2048).shape == (1, 16)
    assert cs.align_codec_rows(None, [], 2048).shape == (0, 0)
    assert cs.align_codec_rows(torch.zeros(0, 16, dtype=torch.long), [], 2048).shape == (0, 16)


def test_session_config_codec_mode_validation():
    cfg = StreamingSpeechSessionConfig(output_mode="codec_tokens", stream_audio=True, response_format="pcm")
    assert cfg.output_mode == "codec_tokens"
    with pytest.raises(ValueError):
        StreamingSpeechSessionConfig(output_mode="codec_tokens", stream_audio=False, response_format="pcm")
    with pytest.raises(ValueError):
        StreamingSpeechSessionConfig(
            output_mode="codec_tokens", stream_audio=True, response_format="pcm", word_timestamps=True
        )
    assert StreamingSpeechSessionConfig().output_mode == "pcm"


def test_codec_start_message_fields():
    spec = {"codebooks": 16, "codebook_size": 2048, "frame_rate_hz": 12.5, "sample_rate": 24000, "decoder_id": "x"}
    msg = cs.codec_start_message(utterance_index=1, sentence_index=0, sentence_text="hi", spec=spec, model="m")
    assert msg["type"] == "codec.start" and msg["codebooks"] == 16 and msg["header_format"] == cs.HEADER_FMT
    assert msg["decoder_id"] == "x" and msg["version"] == cs.CODEC_STREAM_VERSION


# --------------------------------------------------------------------- handler
def _codec_app(mocker, *, supports: bool = True, cancel_after: int | None = None):
    from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
    from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler

    speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
    speech_service.forced_aligner_enabled = False
    speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
    adapter = mocker.MagicMock()
    adapter.supports_codec_stream = supports
    adapter.codec_stream_spec = {
        "codebooks": 16,
        "codebook_size": 2048,
        "frame_rate_hz": 12.5,
        "sample_rate": 24000,
        "decoder_id": "d",
    }
    speech_service._get_tts_adapter = mocker.MagicMock(return_value=adapter)
    speech_service.engine_client = mocker.MagicMock()
    speech_service.engine_client.abort = mocker.AsyncMock()

    async def mock_codec_chunks(_generator, _request_id, *, tts_params=None, min_frames=1):
        import asyncio

        for i in range(3):
            await asyncio.sleep(0.01)
            yield np.full((1 + i, 16), i, dtype=np.int64), i == 2

    speech_service._generate_codec_chunks = mock_codec_chunks
    handler = OmniStreamingSpeechHandler(speech_service=speech_service, idle_timeout=5.0, config_timeout=5.0)
    app = FastAPI()

    @app.websocket("/v1/audio/speech/stream")
    async def ws_endpoint(websocket: WebSocket):
        # NOTE: no ``from __future__ import annotations`` in this module: FastAPI
        # must resolve the WebSocket annotation, or it treats it as a query param.
        await handler.handle_session(websocket)

    return app, speech_service


def test_codec_stream_handler_frames_and_done(mocker):
    from starlette.testclient import TestClient

    app, svc = _codec_app(mocker)
    with TestClient(app) as client, client.websocket_connect("/v1/audio/speech/stream") as ws:
        ws.send_json(
            {"type": "session.config", "output_mode": "codec_tokens", "stream_audio": True, "response_format": "pcm"}
        )
        ws.send_json({"type": "input.text", "text": "hello"})
        ws.send_json({"type": "input.done"})
        start = ws.receive_json()
        assert start["type"] == "codec.start" and start["codebooks"] == 16
        frames = []
        while True:
            msg = ws.receive()
            if "bytes" in msg and msg["bytes"] is not None:
                frames.append(cs.unpack_codec_frame(msg["bytes"], 16))
                continue
            done = __import__("json").loads(msg["text"])
            break
        assert [f[1] for f in frames] == [1, 2, 3]
        assert [f[0] for f in frames] == [0, 1, 2]
        assert frames[-1][2] & cs.FLAG_EOS
        assert done["type"] == "codec.done" and done["total_frames"] == 6 and done["cancelled"] is False
        assert ws.receive_json()["type"] == "session.done"
    svc.engine_client.abort.assert_not_awaited()


def test_codec_stream_rejected_when_adapter_lacks_support(mocker):
    from starlette.testclient import TestClient

    app, _ = _codec_app(mocker, supports=False)
    with TestClient(app) as client, client.websocket_connect("/v1/audio/speech/stream") as ws:
        ws.send_json(
            {"type": "session.config", "output_mode": "codec_tokens", "stream_audio": True, "response_format": "pcm"}
        )
        ws.send_json({"type": "input.text", "text": "hello"})
        ws.send_json({"type": "input.done"})
        err = ws.receive_json()
        assert err["type"] == "error" and "not supported" in err["message"]


def test_input_cancel_mid_stream_aborts_request(mocker):
    from starlette.testclient import TestClient

    app, svc = _codec_app(mocker)
    with TestClient(app) as client, client.websocket_connect("/v1/audio/speech/stream") as ws:
        ws.send_json(
            {"type": "session.config", "output_mode": "codec_tokens", "stream_audio": True, "response_format": "pcm"}
        )
        ws.send_json({"type": "input.text", "text": "hello"})
        ws.send_json({"type": "input.done"})
        assert ws.receive_json()["type"] == "codec.start"
        ws.send_json({"type": "input.cancel"})
        types = []
        while True:
            msg = ws.receive()
            if "bytes" in msg and msg["bytes"] is not None:
                continue
            payload = __import__("json").loads(msg["text"])
            types.append(payload["type"])
            if payload["type"] == "session.done":
                break
            if payload["type"] == "codec.done":
                assert payload["cancelled"] is True
        assert "codec.done" in types and types[-1] == "session.done"
    svc.engine_client.abort.assert_awaited()