# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """The session runner driven by the real PersonaPlex plugin. Same harness as ``test_session_runner.py`` (recording stage port, typed commands through the manager, stage outputs pushed by hand), plugin swapped: this pins the lockstep contract -- one Stage 0 submission per 80 ms frame, prefill slots only on the first append of an epoch, auto-response implied by ``supports_client_commit=False``, cumulative Code2Wav output projected as 24 kHz deltas, and a cancel restarting the epoch. """ from __future__ import annotations import asyncio import base64 import struct from types import SimpleNamespace from typing import Any import numpy as np import pytest from tests.engine.duplex.test_session_runner import ( SESSION_ID, Harness, RecordingStagePort, _fake_encode_audio, close_harness, find, types, ) from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig from vllm_omni.engine.duplex import commands from vllm_omni.engine.duplex.config import DuplexSessionConfig from vllm_omni.engine.duplex.messages import DuplexControlResultMessage, OpenDuplexSessionMessage from vllm_omni.engine.duplex.session.manager import DuplexSessionManager from vllm_omni.model_executor.models.personaplex.duplex import stage0 from vllm_omni.model_executor.models.personaplex.duplex.config import FRAME_SIZE, SAMPLE_RATE from vllm_omni.model_executor.models.personaplex.duplex.plugin import PersonaPlexDuplexPlugin pytestmark = [pytest.mark.core_model, pytest.mark.cpu] PREFILL_SLOTS = 4 @pytest.fixture(autouse=True) def _fake_prefill(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(stage0, "personaplex_prefill_slots", lambda model_path, voice, persona: PREFILL_SLOTS) async def open_personaplex_harness(*, extra_body: dict[str, object] | None = None) -> Harness: plugin = PersonaPlexDuplexPlugin(_fake_encode_audio) port = RecordingStagePort(stage_count=2) output: asyncio.Queue[Any] = asyncio.Queue() results: asyncio.Queue[Any] = asyncio.Queue() manager = DuplexSessionManager( plugin=plugin, stage_port=port, output_sink=output, result_sink=results, runtime_config=DuplexSessionRuntimeConfig(), model_config=SimpleNamespace(model="/models/personaplex-7b-v1"), ) config = DuplexSessionConfig( model="nvidia/personaplex-7b-v1", modalities=["audio", "text"], instructions="Be concise.", voice="NATF2.pt", extra_body=dict(extra_body or {}), ) await manager.handle(OpenDuplexSessionMessage(control_id="c-open", session_id=SESSION_ID, session_config=config)) result = await asyncio.wait_for(results.get(), timeout=2.0) assert isinstance(result, DuplexControlResultMessage) and result.ok, result harness = Harness(manager=manager, port=port, output=output, results=results, runner=manager.runners[SESSION_ID]) await harness.settle() return harness def frame(samples: int = FRAME_SIZE, *, value: float = 0.05) -> commands.AppendAudio: return commands.AppendAudio( audio=struct.pack(f"<{samples}f", *([value] * samples)), format="pcm_f32le", sample_rate_hz=SAMPLE_RATE, is_speech=True, ) def code2wav_output(request_id: str, *, samples: int, text: str) -> SimpleNamespace: """A cumulative Stage 1 (Code2Wav) output the way the orchestrator hands it to the runner.""" return SimpleNamespace( request_id=request_id, finished=False, outputs=[SimpleNamespace(text=text, token_ids=[], multimodal_output={})], multimodal_output={"audio": np.zeros(samples, dtype=np.float32), "sr": SAMPLE_RATE}, ) def submitted_duplex(port: RecordingStagePort) -> list[dict[str, object]]: return [dict(sub.prompt["model_intermediate_buffer"]["duplex"]) for sub in port.submissions] @pytest.mark.asyncio async def test_open_advertises_the_lockstep_capabilities_and_reserves_stage0() -> None: h = await open_personaplex_harness() try: created = find(h.events, "session.created") capabilities = created.session["capabilities"] assert capabilities["chunk_period_ms"] == 80 assert capabilities["supports_client_commit"] is False assert capabilities["supports_barge_in"] is False # Not on the wire: the api server reads it from the plugin to decide the chat route. assert h.session.capabilities.supports_chat_completions is False assert [context.request_id for context in h.port.ensured] == [h.stage0_request_id(epoch=0)] assert h.session.runtime_config["personaplex_prefill_slots"] == PREFILL_SLOTS assert h.session.runtime_config["personaplex_voice_prompt"] == "NATF2.pt" finally: await close_harness(h) @pytest.mark.asyncio async def test_a_session_without_the_auto_response_flag_still_auto_responds() -> None: h = await open_personaplex_harness(extra_body={}) try: assert h.runner.out.auto_responds() is True finally: await close_harness(h) @pytest.mark.asyncio async def test_each_frame_is_one_stage0_submission_with_prefill_only_on_the_first() -> None: h = await open_personaplex_harness() try: await h.run(frame()) await h.run(frame()) await h.run(frame()) assert len(h.port.submissions) == 3 assert [len(sub.prompt["prompt_token_ids"]) for sub in h.port.submissions] == [1 + PREFILL_SLOTS, 1, 1] duplex = submitted_duplex(h.port) assert [item["seq"] for item in duplex] == [1, 2, 3] assert all(item["epoch"] == 0 for item in duplex) assert all(item["session_id"] == SESSION_ID for item in duplex) assert all("incarnation" not in item for item in duplex) assert all(sub.context.request_id == h.stage0_request_id(epoch=0) for sub in h.port.submissions) assert [sub.already_submitted for sub in h.port.submissions] == [False, True, True] assert h.port.submissions[0].context.stage_sampling_params.max_tokens == 1 assert not [event for event in h.events if event.type == "error"] finally: await close_harness(h) @pytest.mark.asyncio async def test_half_frames_are_buffered_until_a_whole_frame_exists() -> None: h = await open_personaplex_harness() try: await h.run(frame(FRAME_SIZE // 2)) assert h.port.submissions == [] assert h.runner.model_state.audio_buffer.pending_byte_count == FRAME_SIZE * 2 await h.run(frame(FRAME_SIZE // 2)) assert len(h.port.submissions) == 1 assert h.runner.model_state.audio_buffer.pending_byte_count == 0 payload = submitted_duplex(h.port)[0]["payload"] assert isinstance(payload, dict) assert len(base64.b64decode(payload["audio"])) == FRAME_SIZE * 4 finally: await close_harness(h) @pytest.mark.asyncio async def test_a_wrong_rate_append_is_refused_without_a_submission() -> None: h = await open_personaplex_harness() try: events = await h.run( commands.AppendAudio(audio=b"\x00" * 64, format="pcm_f32le", sample_rate_hz=16000, is_speech=None) ) error = find(events, "error") assert error.code == "bad_event" assert "24000" in error.message assert h.port.submissions == [] finally: await close_harness(h) @pytest.mark.asyncio async def test_cumulative_code2wav_output_becomes_24k_audio_and_text_deltas() -> None: h = await open_personaplex_harness() try: await h.run(frame()) request_id = h.stage0_request_id(epoch=0) first = await h.deliver_and_settle(code2wav_output(request_id, samples=FRAME_SIZE, text="he"), stage_id=1) second = await h.deliver_and_settle( code2wav_output(request_id, samples=2 * FRAME_SIZE, text="hello"), stage_id=1 ) assert types(first)[0] == "response.created" delta = find(first, "response.output_audio.delta") assert delta.delta == f"wav-{FRAME_SIZE}" assert delta.sample_rate_hz == SAMPLE_RATE assert find(first, "response.output_audio_transcript.delta").delta == "he" assert h.session.active_response_id is not None [delta_2] = [event for event in second if event.type == "response.output_audio.delta"] assert delta_2.delta == f"wav-{FRAME_SIZE}" assert find(second, "response.output_audio_transcript.delta").delta == "llo" assert not [event for event in second if event.type == "response.created"] finally: await close_harness(h) @pytest.mark.asyncio async def test_cancel_restarts_the_epoch_and_the_next_frame_replays_the_prefill() -> None: h = await open_personaplex_harness() try: await h.run(frame()) old_request_id = h.stage0_request_id(epoch=0) await h.deliver_and_settle(code2wav_output(old_request_id, samples=FRAME_SIZE, text="he"), stage_id=1) cancelled = await h.run(commands.CancelResponse()) await h.run(frame()) assert h.session.epoch == 1 assert "audio.cancelled" in types(cancelled) or "response.done" in types(cancelled) assert ([old_request_id], True) in h.port.cleanups assert h.port.submissions[-1].context.request_id == h.stage0_request_id(epoch=1) assert submitted_duplex(h.port)[-1]["seq"] == 1 assert len(h.port.submissions[-1].prompt["prompt_token_ids"]) == 1 + PREFILL_SLOTS finally: await close_harness(h) @pytest.mark.asyncio async def test_close_aborts_the_stage0_request() -> None: h = await open_personaplex_harness() await h.run(frame()) request_id = h.stage0_request_id(epoch=0) await h.run(commands.CloseSession(reason="client_close")) assert h.port.cleanups == [([request_id], True)] await close_harness(h)