vllm-omni/tests/entrypoints/test_duplex_omni.py

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

"""``DuplexOmni`` / ``DuplexSessionHandle``: the thin Python API over a fake engine."""

from __future__ import annotations

import asyncio
from dataclasses import replace
from typing import Any

import pytest

from tests.entrypoints.test_omni_entrypoints import FakeAsyncOmniEngine
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.engine.duplex import commands
from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
from vllm_omni.engine.duplex.events import (
    AudioDelta,
    DuplexEvent,
    SessionClosed,
    SessionCreated,
    SessionExpired,
)
from vllm_omni.engine.duplex.messages import (
    DuplexControlResultMessage,
    DuplexSessionError,
    DuplexSessionEventMessage,
)
from vllm_omni.engine.messages import ErrorMessage
from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle

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


class FakeDuplexEngine(FakeAsyncOmniEngine):
    """The engine surface ``DuplexOmni`` uses; every call is recorded."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.duplex_session_config = DuplexSessionRuntimeConfig(max_sessions=2)
        self.duplex_capabilities = DuplexCapabilities(supports_session_resume=True)
        self.opened: list[tuple[str, DuplexSessionConfig]] = []
        self.outputs: dict[str, DuplexOutputBuffer] = {}
        self.closed: list[tuple[str, str]] = []
        self.resumed: list[tuple[str, int]] = []
        #: The control id of every resume RPC, in order (a replay repeats one).
        self.resume_control_ids: list[str | None] = []
        #: Per session: the control id that produced the current lease generation, and that generation.
        self.lease: dict[str, tuple[str | None, int]] = {}
        #: How many resume RPCs, after applying the resume, answer with a
        #: timeout: the engine landed it but the waiter never saw the result.
        self.resume_timeouts_remaining = 0
        self.touched: list[tuple[str, str]] = []
        #: The lease-generation fence of every detach, in order (``None`` = unfenced).
        self.detach_generations: list[int | None] = []
        #: When set, ``resume_session_async`` parks on it after signalling ``resume_started``:
        #: the engine has applied the resume, its result has not reached the caller.
        self.resume_gate: asyncio.Event | None = None
        self.resume_started = asyncio.Event()
        self.commands: list[tuple[str, commands.DuplexCommand]] = []
        self.open_error: DuplexSessionError | None = None
        self.close_error: DuplexSessionError | None = None

    def _result(self, operation: str, session_id: str, **fields: Any) -> DuplexControlResultMessage:
        return DuplexControlResultMessage(control_id="c", operation=operation, session_id=session_id, ok=True, **fields)

    async def open_session_async(self, session_id, session_config, *, output_buffer, timeout=None):
        self.opened.append((session_id, session_config))
        if self.open_error is not None:
            raise self.open_error
        self.outputs[session_id] = output_buffer
        return self._result(
            "open",
            session_id,
            capabilities=self.duplex_capabilities,
            public_session={"id": session_id, "idle_timeout_s": 7},
            lease_generation=0,
        )

    async def close_session_async(self, session_id, *, reason="client_close", timeout=None):
        self.closed.append((session_id, reason))
        if self.close_error is not None:
            raise self.close_error
        # The runner emits session.closed before answering the RPC.
        self.emit(session_id, SessionClosed(reason=reason))
        return self._result("close", session_id)

    async def resume_session_async(self, session_id, *, expected_lease_generation, control_id=None, timeout=None):
        self.resumed.append((session_id, expected_lease_generation))
        self.resume_control_ids.append(control_id)
        resumed_by, generation = self.lease.get(session_id, (None, 0))
        if control_id is not None and control_id == resumed_by:
            # The engine's replay answer: the generation this resume produced.
            # Its answer can be lost to a timeout exactly like the first one.
            if self.resume_timeouts_remaining > 0:
                self.resume_timeouts_remaining -= 1
                raise DuplexSessionError("duplex resume timed out", code="timeout", retryable=True)
            return self._result("resume", session_id, lease_generation=generation)
        if expected_lease_generation != generation:
            raise DuplexSessionError("duplex lease generation mismatch", code="session_resume_conflict")
        generation += 1
        self.lease[session_id] = (control_id, generation)
        if self.resume_gate is not None:
            self.resume_started.set()
            await self.resume_gate.wait()
        if self.resume_timeouts_remaining > 0:
            self.resume_timeouts_remaining -= 1
            raise DuplexSessionError("duplex resume timed out", code="timeout", retryable=True)
        return self._result("resume", session_id, lease_generation=generation)

    async def touch_session_async(self, session_id, *, activity, expected_lease_generation=None, timeout=None):
        self.touched.append((session_id, activity))
        if activity == "detach":
            self.detach_generations.append(expected_lease_generation)
        return self._result("touch", session_id)

    async def submit_command_async(self, session_id, command):
        self.commands.append((session_id, command))

    def emit(self, session_id: str, event: DuplexEvent) -> None:
        # The session manager binds the session identity to every event it emits.
        event = replace(event, session_id=session_id)
        if isinstance(event, SessionClosed):
            self.outputs.pop(session_id, None)
            self.output_q.put(DuplexSessionEventMessage(session_id=session_id, event=event))
        elif session_id in self.outputs:
            self.outputs[session_id].put(event)


def _make_omni(monkeypatch: pytest.MonkeyPatch) -> tuple[DuplexOmni, FakeDuplexEngine]:
    engine = FakeDuplexEngine()
    monkeypatch.setattr("vllm_omni.entrypoints.duplex_omni.DuplexOmniEngine", lambda *args, **kwargs: engine)
    monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
    return DuplexOmni("dummy-model"), engine


async def _collect(handle: DuplexSessionHandle, count: int, *, timeout_s: float = 2.0) -> list[DuplexEvent]:
    events: list[DuplexEvent] = []

    async def consume() -> None:
        async for event in handle.events():
            events.append(event)
            if len(events) >= count:
                break

    await asyncio.wait_for(consume(), timeout=timeout_s)
    return events


@pytest.mark.asyncio
async def test_open_session_allocates_the_id_and_ignores_client_ids(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session({"session_id": "mine", "id": "mine-too", "instructions": "hi"})

        assert handle.session_id.startswith("duplex-") and len(handle.session_id) == len("duplex-") + 32
        assert handle.session_id != "mine"
        ((session_id, config),) = engine.opened
        assert session_id == handle.session_id
        assert engine.outputs[session_id] is handle._outbox
        assert isinstance(config, DuplexSessionConfig)
        assert config.model == "dummy-model" and config.instructions == "hi"
        assert handle.capabilities is engine.duplex_capabilities
        assert handle.public_session["id"] == handle.session_id
        assert omni.sessions == {handle.session_id: handle}
        assert omni.get_session(handle.session_id) is handle
        assert omni.active_session_count() == 1

        second = await omni.open_session(DuplexSessionConfig(instructions="typed"))
        assert second.session_id != handle.session_id
        assert engine.opened[1][1].model == "dummy-model"
        assert omni.duplex_session_config.max_sessions == 2
        assert omni.duplex_capabilities.supports_session_resume is True
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_events_skips_audio_invalidated_after_dequeue(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        output = engine.outputs[handle.session_id]
        original_get = output.get

        async def get_then_invalidate() -> DuplexEvent | None:
            event = await original_get()
            if isinstance(event, AudioDelta) and event.response_id == "cancelled":
                output.invalidate("cancelled", through_epoch=0)
            return event

        monkeypatch.setattr(output, "get", get_then_invalidate)
        engine.emit(handle.session_id, AudioDelta(response_id="cancelled", epoch=0, delta="old"))
        engine.emit(handle.session_id, AudioDelta(response_id="current", epoch=1, delta="new"))

        events = await _collect(handle, 1)

        assert isinstance(events[0], AudioDelta)
        assert events[0].response_id == "current" and events[0].delta == "new"
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_open_failure_drops_the_pending_handle_and_raises_the_engine_code(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    engine.open_error = DuplexSessionError("full", code="resource_exhausted", retryable=True)
    try:
        with pytest.raises(DuplexSessionError) as excinfo:
            await omni.open_session()
        assert excinfo.value.code == "resource_exhausted"
        assert omni.sessions == {}
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_events_are_routed_to_their_handle_and_end_on_session_closed(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        engine.emit("duplex-unknown", AudioDelta(delta="zzz"))  # dropped
        engine.emit(handle.session_id, SessionCreated(session={"id": handle.session_id}))
        engine.emit(handle.session_id, AudioDelta(response_id="r1", delta="abc"))
        engine.emit(handle.session_id, SessionClosed(reason="client_close"))

        events = await _collect(handle, 10)

        assert [event.type for event in events] == ["session.created", "response.output_audio.delta", "session.closed"]
        assert events[0].session_id == handle.session_id
        assert handle.closed and handle.close_reason == "client_close"
        assert await handle.wait_closed() == "client_close"
        assert handle.session_id not in omni.sessions
        with pytest.raises(DuplexSessionError) as excinfo:
            await handle.submit(commands.Heartbeat())
        assert excinfo.value.code == "session_closed"
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_events_is_single_consumer_but_may_be_reentered(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        engine.emit(handle.session_id, AudioDelta(delta="one"))
        first = handle.events()
        assert (await first.__anext__()).type == "response.output_audio.delta"
        with pytest.raises(RuntimeError, match="active events"):
            await handle.events().__anext__()
        await first.aclose()

        engine.emit(handle.session_id, SessionExpired(reason="idle_ttl_expired"))
        events = await _collect(handle, 1)
        assert events[0].type == "session.expired" and events[0].is_terminal
        assert handle.closed and handle.close_reason == "idle_ttl_expired"
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_close_session_waits_for_the_typed_session_closed_event(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        consumer = asyncio.create_task(_collect(handle, 1))
        await handle.close(reason="done")

        assert engine.closed == [(handle.session_id, "done")]
        assert handle.closed and handle.close_reason == "done"
        assert (await consumer)[0].type == "session.closed"
        # Idempotent: a second close is a no-op.
        await omni.close_session(handle.session_id)
        assert len(engine.closed) == 1
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_close_failure_marks_the_handle_closed_and_propagates(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    engine.close_error = DuplexSessionError("boom", code="engine_error")
    try:
        handle = await omni.open_session()
        with pytest.raises(DuplexSessionError, match="boom"):
            await omni.close_session(handle.session_id)
        assert handle.closed
        # A failed close must not leave a dead handle shadowing the id.
        assert omni.sessions == {}
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_engine_death_ends_every_session_handle(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        engine.output_q.put(ErrorMessage(error="orchestrator died", fatal=True))

        await asyncio.wait_for(handle.wait_closed(), timeout=2.0)
        assert handle.closed
        assert "engine_dead" in (handle.close_reason or "")
        assert omni.sessions == {}
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_resume_detach_and_touch_forward_to_the_engine(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        resumed = await omni.resume_session(handle.session_id, expected_lease_generation=0)
        assert resumed is handle and handle.lease_generation == 1
        assert engine.resumed == [(handle.session_id, 0)]

        await omni.detach_session(handle.session_id)
        await omni.touch_session(handle.session_id, activity="heartbeat")
        assert engine.touched == [(handle.session_id, "detach"), (handle.session_id, "heartbeat")]

        for call in (
            omni.resume_session("duplex-nope", expected_lease_generation=0),
            omni.detach_session("duplex-nope"),
        ):
            with pytest.raises(DuplexSessionError) as excinfo:
                await call
            assert excinfo.value.code == "unknown_session"
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_handle_wrappers_build_typed_commands_in_caller_order(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        await handle.append_audio(b"\x00\x01", sample_rate_hz=16000, is_speech=True)
        await handle.commit(create_response=True)
        await handle.create_response({"instructions": "x"})
        await handle.cancel_response("resp-1")
        await handle.barge_in()
        await handle.ack_playback(120.7, response_id="resp-1", committed_ms=100)
        await handle.update({"temperature": 0.2})
        await handle.create_item({"id": "item_1", "type": "message"}, previous_item_id="root")
        await handle.delete_item("item_1")
        await handle.truncate_item("item_1", audio_end_ms=50)
        await handle.signal_turn("user_started", {"x": 1})
        await handle.heartbeat()

        submitted = [command for session_id, command in engine.commands if session_id == handle.session_id]
        assert [type(command) for command in submitted] == [
            commands.AppendAudio,
            commands.Commit,
            commands.CreateResponse,
            commands.CancelResponse,
            commands.BargeIn,
            commands.AckPlayback,
            commands.UpdateSession,
            commands.CreateItem,
            commands.DeleteItem,
            commands.TruncateItem,
            commands.SignalTurn,
            commands.Heartbeat,
        ]
        append, commit, create, cancel, _, ack, update, item, delete, truncate, signal, _ = submitted
        assert append.audio == b"\x00\x01" and append.sample_rate_hz == 16000 and append.is_speech is True
        assert commit.create_response is True and commit.final is True
        assert create.options == {"instructions": "x"}
        assert cancel.response_id == "resp-1"
        assert ack.played_ms == 120 and ack.committed_ms == 100 and ack.response_id == "resp-1"
        assert update.patch == {"temperature": 0.2}
        assert item.previous_item_id == "root" and item.item["id"] == "item_1"
        assert delete.item_id == "item_1"
        assert truncate.audio_end_ms == 50
        assert signal.event == "user_started" and signal.signal_payload == {"x": 1}
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_shutdown_closes_every_handle(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    first = await omni.open_session()
    second = await omni.open_session()

    omni.shutdown()

    assert first.closed and second.closed
    assert first.close_reason == "shutdown"
    assert omni.sessions == {}
    assert engine.shutdown_called


@pytest.mark.asyncio
async def test_detach_session_forwards_the_lease_generation_fence(monkeypatch) -> None:
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        await omni.detach_session(handle.session_id, expected_lease_generation=3)
        await omni.detach_session(handle.session_id)
        assert engine.detach_generations == [3, None]
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_a_resume_cancelled_mid_rpc_is_adopted_and_detached_once_it_lands(monkeypatch) -> None:
    """The resume RPC runs in an executor: cancelling its waiter does not stop the engine.

    Without compensation the engine ends up with the lease resumed (no
    disconnect grace) for a connection that is gone, and the handle keeps the
    old generation, so a later reconnect resumes against a stale lease. The
    outcome has to be observed off the cancelled task: adopt the generation
    the resume produced and detach exactly that lease.
    """
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        engine.resume_gate = asyncio.Event()
        resume = asyncio.create_task(omni.resume_session(handle.session_id, expected_lease_generation=0))
        await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0)

        resume.cancel()
        with pytest.raises(asyncio.CancelledError):
            await resume
        # Nothing to give back yet: the engine has not answered.
        assert handle.lease_generation == 0
        assert engine.detach_generations == []

        engine.resume_gate.set()
        await asyncio.gather(*list(omni._resume_compensations))
        assert handle.lease_generation == 1, "the generation the resume produced is adopted"
        assert engine.touched == [(handle.session_id, "detach")]
        assert engine.detach_generations == [1], "the detach is fenced on that generation"
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_a_resume_cancelled_then_timed_out_is_settled_by_replaying_it(monkeypatch) -> None:
    """A timeout after the cancel leaves the outcome unknown: the resume is replayed under its id.

    The engine may have applied the resume after the waiter gave up, and the
    RPC router drops the late answer. Replaying the same control id makes the
    engine report the generation that resume produced (or land it now), so
    the handle adopts the real generation and the lease is detached again.
    """
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        engine.resume_gate = asyncio.Event()
        engine.resume_timeouts_remaining = 1
        resume = asyncio.create_task(omni.resume_session(handle.session_id, expected_lease_generation=0))
        await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0)

        resume.cancel()
        with pytest.raises(asyncio.CancelledError):
            await resume
        engine.resume_gate.set()
        await asyncio.gather(*list(omni._resume_compensations))

        assert len(engine.resume_control_ids) == 2, "the resume was replayed once"
        assert engine.resume_control_ids[0] == engine.resume_control_ids[1] is not None
        assert engine.lease[handle.session_id][1] == 1, "the replay did not resume a second time"
        assert handle.lease_generation == 1
        assert engine.detach_generations == [1]
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_an_abandoned_resume_is_settled_through_the_callers_callback(monkeypatch) -> None:
    """A caller that owns attachments decides what the landed generation means; DuplexOmni does not detach."""
    omni, engine = _make_omni(monkeypatch)
    try:
        handle = await omni.open_session()
        settled: list[int] = []

        async def on_abandoned(lease_generation: int) -> None:
            settled.append(lease_generation)

        engine.resume_gate = asyncio.Event()
        resume = asyncio.create_task(
            omni.resume_session(handle.session_id, expected_lease_generation=0, on_abandoned=on_abandoned)
        )
        await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0)
        resume.cancel()
        with pytest.raises(asyncio.CancelledError):
            await resume
        engine.resume_gate.set()
        await asyncio.gather(*list(omni._resume_compensations))

        assert settled == [1]
        assert handle.lease_generation == 1
        assert engine.detach_generations == [], "the callback owns the decision"
    finally:
        omni.shutdown()


@pytest.mark.asyncio
async def test_an_abandoned_resume_keeps_replaying_across_repeated_timeouts(monkeypatch) -> None:
    """A replay can time out too; giving up then would reopen the lost-answer window.

    The engine applied the resume, the abandoned RPC and the first replay both
    lose their answers to timeouts, the second replay is answered: the handle
    adopts the generation and the lease is settled, with the lease bumped once.
    """
    omni, engine = _make_omni(monkeypatch)
    try:
        omni._resume_replay_backoff_s = (0.01, 0.02)
        handle = await omni.open_session()
        engine.resume_gate = asyncio.Event()
        engine.resume_timeouts_remaining = 2
        resume = asyncio.create_task(omni.resume_session(handle.session_id, expected_lease_generation=0))
        await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0)

        resume.cancel()
        with pytest.raises(asyncio.CancelledError):
            await resume
        engine.resume_gate.set()
        await asyncio.wait_for(asyncio.gather(*list(omni._resume_compensations)), timeout=5.0)

        assert len(engine.resume_control_ids) == 3, "the abandoned RPC, a replay that timed out, the replay answered"
        assert len(set(engine.resume_control_ids)) == 1
        assert engine.lease[handle.session_id][1] == 1, "the lease was bumped once"
        assert handle.lease_generation == 1
        assert engine.detach_generations == [1]
    finally:
        omni.shutdown()