vllm-omni/vllm_omni/entrypoints/duplex_omni.py

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

"""DuplexOmni: the Python API for full-duplex models.

Sessions run inside the engine (``DuplexSessionRunner`` on the orchestrator
loop of ``DuplexOrchestrator``). This class opens / resumes / closes them and
pipes typed ``DuplexCommand`` objects in and typed ``DuplexEvent`` objects out
through a ``DuplexSessionHandle``. It holds no session state beyond the
handle registry, so the websocket handler and ``InlineDuplexClient`` are both
thin consumers of the same surface.

Example::

    omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True)
    async with await omni.open_session({"modalities": ["audio", "text"]}) as session:
        async def consume():
            async for event in session.events():
                if isinstance(event, AudioDelta):
                    play(event.audio)
        task = asyncio.create_task(consume())
        await session.submit(AppendAudio(audio=pcm_bytes, format="pcm16", sample_rate_hz=16000))
        await session.submit(Commit())
"""

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator, Mapping, Sequence
from contextlib import suppress
from typing import Any
from uuid import uuid4

from vllm.logger import init_logger

from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.engine.duplex import commands as duplex_commands
from vllm_omni.engine.duplex.commands import DuplexCommand
from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig, ResponseCreateOptions
from vllm_omni.engine.duplex.events import DuplexEvent, SessionClosed
from vllm_omni.engine.duplex.messages import (
    DuplexControlResultMessage,
    DuplexSessionError,
    DuplexSessionEventMessage,
)
from vllm_omni.engine.duplex_omni_engine import DuplexOmniEngine
from vllm_omni.entrypoints.async_omni import AsyncOmni

logger = init_logger(__name__)

_DEFAULT_CONTROL_TIMEOUT_S = 10.0


class DuplexSessionHandle:
    """Client-side view of one engine-resident duplex session.

    Connection-independent: a websocket may detach and a later one resume the
    same handle. ``events()`` is single-consumer at any given time but may be
    re-entered after the previous iterator was closed (resume).
    """

    def __init__(self, omni: DuplexOmni, session_id: str) -> None:
        self._omni = omni
        #: Allocated by ``DuplexOmni.open_session``; unique for the engine's lifetime.
        self.session_id = session_id
        self.capabilities: DuplexCapabilities = DuplexCapabilities()
        self.public_session: dict[str, object] = {}
        self.lease_generation: int = 0
        self._outbox: asyncio.Queue[DuplexEvent | None] = asyncio.Queue()
        self._closed = False
        self._close_reason: str | None = None
        self._closed_event = asyncio.Event()
        self._consumer_active = False

    # ---- state ----

    @property
    def closed(self) -> bool:
        return self._closed

    @property
    def close_reason(self) -> str | None:
        return self._close_reason

    def __repr__(self) -> str:  # pragma: no cover - debugging aid
        return f"DuplexSessionHandle(session_id={self.session_id!r}, closed={self._closed})"

    # ---- input ----

    async def submit(self, command: DuplexCommand) -> None:
        """Enqueue one command in caller order; rejections arrive as ``ErrorEvent`` on ``events()``."""
        if self._closed:
            raise DuplexSessionError(
                f"duplex session {self.session_id} is closed", code="session_closed", session_id=self.session_id
            )
        await self._omni.engine.submit_command_async(self.session_id, command)

    async def append_audio(
        self,
        audio: bytes,
        *,
        format: str = "pcm16",
        sample_rate_hz: int | None = None,
        is_speech: bool | None = None,
        video_frames: Sequence[str] | None = None,
        duration_ms: int | None = None,
        audio_end_ms: int | None = None,
        hints: Mapping[str, object] | None = None,
    ) -> None:
        await self.submit(
            duplex_commands.AppendAudio(
                audio=audio,
                format=format,
                sample_rate_hz=sample_rate_hz,
                is_speech=is_speech,
                video_frames=tuple(video_frames or ()),
                duration_ms=duration_ms,
                audio_end_ms=audio_end_ms,
                hints=dict(hints or {}),
            )
        )

    async def append_text(self, text: str) -> None:
        await self.submit(duplex_commands.AppendText(text=text))

    async def commit(
        self,
        *,
        final: bool = True,
        create_response: bool | None = None,
        is_speech: bool | None = None,
    ) -> None:
        await self.submit(duplex_commands.Commit(final=final, create_response=create_response, is_speech=is_speech))

    async def create_response(self, options: ResponseCreateOptions | Mapping[str, object] | None = None) -> None:
        if isinstance(options, ResponseCreateOptions):
            payload: dict[str, object] = {
                key: value
                for key, value in (
                    ("instructions", options.instructions),
                    ("voice", options.voice),
                    ("output_audio_format", options.response_format),
                    ("temperature", options.temperature),
                    ("max_output_tokens", options.max_tokens),
                    ("speed", options.speed),
                    ("modalities", list(options.modalities) if options.modalities is not None else None),
                )
                if value is not None
            }
            if options.extra_body:
                payload["extra_body"] = dict(options.extra_body)
        else:
            payload = dict(options or {})
        await self.submit(duplex_commands.CreateResponse(options=payload))

    async def clear_input(self) -> None:
        await self.submit(duplex_commands.ClearInput())

    async def cancel_input(self) -> None:
        await self.submit(duplex_commands.CancelInput())

    async def cancel_response(self, response_id: str | None = None) -> None:
        await self.submit(duplex_commands.CancelResponse(response_id=response_id))

    async def barge_in(self) -> None:
        await self.submit(duplex_commands.BargeIn())

    async def clear_output_audio(self, response_id: str | None = None) -> None:
        await self.submit(duplex_commands.ClearOutputAudio(response_id=response_id))

    async def signal_turn(self, event: str, payload: Mapping[str, object] | None = None) -> None:
        await self.submit(duplex_commands.SignalTurn(event=event, signal_payload=dict(payload or {})))

    async def update(self, session_patch: Mapping[str, object]) -> None:
        await self.submit(duplex_commands.UpdateSession(patch=dict(session_patch)))

    async def ack_playback(
        self,
        played_ms: int,
        *,
        response_id: str | None = None,
        item_id: str | None = None,
        committed_ms: int | None = None,
    ) -> None:
        await self.submit(
            duplex_commands.AckPlayback(
                played_ms=int(played_ms),
                committed_ms=int(committed_ms) if committed_ms is not None else None,
                response_id=response_id,
                item_id=item_id,
            )
        )

    async def heartbeat(self) -> None:
        await self.submit(duplex_commands.Heartbeat())

    async def create_item(self, item: Mapping[str, object], *, previous_item_id: str | None = None) -> None:
        await self.submit(duplex_commands.CreateItem(item=dict(item), previous_item_id=previous_item_id))

    async def delete_item(self, item_id: str) -> None:
        await self.submit(duplex_commands.DeleteItem(item_id=item_id))

    async def truncate_item(self, item_id: str, *, audio_end_ms: int, content_index: int = 0) -> None:
        await self.submit(
            duplex_commands.TruncateItem(item_id=item_id, audio_end_ms=int(audio_end_ms), content_index=content_index)
        )

    async def close(self, *, reason: str = "client_close", timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S) -> None:
        if self._closed:
            return
        await self._omni.close_session(self.session_id, reason=reason, timeout=timeout)

    # ---- output ----

    async def events(self) -> AsyncIterator[DuplexEvent]:
        """Ordered public events; ends after ``session.closed`` / ``session.expired``."""
        if self._consumer_active:
            raise RuntimeError(f"duplex session {self.session_id} already has an active events() consumer")
        self._consumer_active = True
        try:
            while True:
                if self._closed and self._outbox.empty():
                    return
                event = await self._outbox.get()
                if event is None:
                    if self._closed and self._outbox.empty():
                        return
                    continue
                yield event
                if event.is_terminal:
                    return
        finally:
            self._consumer_active = False

    async def wait_closed(self) -> str:
        await self._closed_event.wait()
        return self._close_reason or "closed"

    async def __aenter__(self) -> DuplexSessionHandle:
        return self

    async def __aexit__(self, exc_type, exc, traceback) -> None:
        try:
            await self.close()
        except DuplexSessionError:
            pass

    # ---- internals used by DuplexOmni ----

    def _adopt(self, result: DuplexControlResultMessage) -> None:
        if result.capabilities is not None:
            self.capabilities = result.capabilities
        if result.public_session:
            self.public_session = dict(result.public_session)
        if result.lease_generation is not None:
            self.lease_generation = int(result.lease_generation)

    def _deliver(self, event: DuplexEvent) -> None:
        self._outbox.put_nowait(event)
        if isinstance(event, SessionClosed):  # SessionExpired is a SessionClosed
            self._mark_closed(event.reason or event.type)

    def _mark_closed(self, reason: str) -> None:
        if self._closed:
            return
        self._closed = True
        self._close_reason = reason
        self._closed_event.set()
        self._outbox.put_nowait(None)


class DuplexOmni(AsyncOmni):
    """Async Python API for full-duplex models (see module docstring).

    Construct it like ``AsyncOmni``. The pipeline must declare
    ``duplex_plugin`` and the deploy config ``session_mode: duplex``; the
    engine raises at startup otherwise. There is no ``generate()``: duplex
    serving is session-only.
    """

    engine: DuplexOmniEngine

    def _create_engine(self, **engine_kwargs: Any) -> DuplexOmniEngine:
        from vllm_omni.entrypoints.duplex.audio_encoding import encode_audio

        return DuplexOmniEngine(duplex_audio_encoder=encode_audio, **engine_kwargs)

    def __init__(self, model: str = "", *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, model=model, **kwargs)
        self._handles: dict[str, DuplexSessionHandle] = {}

    # ---- deployment facts ----

    @property
    def duplex_session_config(self) -> DuplexSessionRuntimeConfig:
        return self.engine.duplex_session_config

    @property
    def duplex_capabilities(self) -> DuplexCapabilities:
        return self.engine.duplex_capabilities

    @property
    def sessions(self) -> Mapping[str, DuplexSessionHandle]:
        return dict(self._handles)

    def active_session_count(self) -> int:
        return sum(1 for handle in self._handles.values() if not handle.closed)

    # ---- session lifecycle ----

    @staticmethod
    def _resolve_session_config(
        config: DuplexSessionConfig | Mapping[str, object] | None,
        *,
        model: str,
    ) -> DuplexSessionConfig:
        """Build the normalized session config the engine will own (the one parse/normalize site).

        A ``session_id`` / ``id`` key inside a Realtime session object is
        ignored: Realtime clients echo the session object back, and the server
        allocates every session id.
        """
        if config is None:
            resolved = DuplexSessionConfig(model=model)
        elif isinstance(config, DuplexSessionConfig):
            resolved = config
        elif isinstance(config, Mapping):
            resolved = DuplexSessionConfig.from_realtime(config, model=model)
        else:
            raise TypeError(f"unsupported duplex session config: {type(config).__name__}")
        if resolved.model is None:
            resolved.model = model
        return resolved.normalized()

    async def open_session(
        self,
        config: DuplexSessionConfig | Mapping[str, object] | None = None,
        *,
        timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
    ) -> DuplexSessionHandle:
        """Open an engine-resident session and return its handle.

        The session id is always allocated here (``duplex-<uuid4 hex>``; never
        reused, so the id alone identifies a session) and the handle is
        registered before the open RPC, so the first event (``session.created``,
        which announces the id) can never arrive before it exists.
        """
        session_id = f"duplex-{uuid4().hex}"
        session_config = self._resolve_session_config(config, model=self.model)
        handle = DuplexSessionHandle(self, session_id)
        self._handles[session_id] = handle
        self._final_output_handler()
        try:
            result = await self.engine.open_session_async(session_id, session_config, timeout=timeout)
        except BaseException:
            if self._handles.get(session_id) is handle:
                self._handles.pop(session_id, None)
            # Dropping the handle only forgets the id here. The engine may still
            # be admitting this session -- a control timeout unregisters the RPC
            # waiter but cannot cancel the manager's open -- so a late admission
            # would hold an admission slot and a Stage0 reservation until idle
            # expiry. Send the compensating close; the manager applies it
            # immediately if the session already landed, and records it for the
            # open to honour on arrival if not.
            with suppress(Exception):
                await self.engine.close_session_async(session_id, reason="open_abandoned", timeout=timeout)
            raise
        handle._adopt(result)
        return handle

    def get_session(self, session_id: str) -> DuplexSessionHandle | None:
        return self._handles.get(session_id)

    async def resume_session(
        self,
        session_id: str,
        *,
        expected_lease_generation: int,
        timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
    ) -> DuplexSessionHandle:
        """Engine lease resume (CAS on the lease generation); returns the existing handle."""
        handle = self._require_handle(session_id)
        result = await self.engine.resume_session_async(
            session_id,
            expected_lease_generation=expected_lease_generation,
            timeout=timeout,
        )
        # The result carries the engine's current public session (state,
        # epoch, turn ...), which a resumed client must see, not the open-time snapshot.
        handle._adopt(result)
        return handle

    async def touch_session(
        self,
        session_id: str,
        *,
        activity: str = "heartbeat",
        timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
    ) -> None:
        self._require_handle(session_id)
        await self.engine.touch_session_async(session_id, activity=activity, timeout=timeout)

    async def detach_session(self, session_id: str, *, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S) -> None:
        """Start the engine-owned disconnect grace; expiry arrives as ``session.expired``."""
        await self.touch_session(session_id, activity="detach", timeout=timeout)

    async def close_session(
        self,
        session_id: str,
        *,
        reason: str = "client_close",
        timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
    ) -> None:
        handle = self._handles.get(session_id)
        if handle is None or handle.closed:
            return
        try:
            await self.engine.close_session_async(session_id, reason=reason, timeout=timeout)
        except Exception:
            # The engine refused or lost the close: the handle is unusable
            # either way, and it must not shadow the id in ``sessions``.
            handle._mark_closed(reason)
            self._handles.pop(session_id, None)
            raise
        # The manager emits session.closed after the stage cleanup, but the
        # event travels on the output queue: wait for it so ``events()`` ends
        # with the typed SessionClosed rather than an abrupt stop.
        try:
            await asyncio.wait_for(handle.wait_closed(), timeout=timeout)
        except asyncio.TimeoutError:
            logger.warning("[DuplexOmni] session.closed for %s did not arrive in time", session_id)
            handle._mark_closed(reason)
            self._handles.pop(session_id, None)

    async def close_all_sessions(self, *, reason: str = "shutdown") -> None:
        for session_id in list(self._handles):
            try:
                await self.close_session(session_id, reason=reason)
            except Exception:
                logger.exception("[DuplexOmni] failed to close session %s", session_id)

    def _require_handle(self, session_id: str) -> DuplexSessionHandle:
        handle = self._handles.get(session_id)
        if handle is None or handle.closed:
            raise DuplexSessionError(
                f"unknown or closed duplex session: {session_id}", code="unknown_session", session_id=session_id
            )
        return handle

    # ---- engine output routing ----

    def _route_engine_message(self, msg: object) -> bool:
        if not isinstance(msg, DuplexSessionEventMessage):
            return False
        handle = self._handles.get(msg.session_id)
        if handle is None:
            logger.debug("[DuplexOmni] dropping event for unknown session %s", msg.session_id)
            return True
        handle._deliver(msg.event)
        if handle.closed:
            self._handles.pop(msg.session_id, None)
        return True

    def _on_engine_dead(self, error: str) -> None:
        for handle in list(self._handles.values()):
            handle._mark_closed(f"engine_dead: {error}")
        self._handles.clear()

    def shutdown(self, timeout: float | None = None) -> None:
        for handle in list(self._handles.values()):
            handle._mark_closed("shutdown")
        self._handles.clear()
        super().shutdown(timeout)


__all__ = ["DuplexOmni", "DuplexSessionHandle"]