544 lines
25 KiB
Python
544 lines
25 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Thin websocket serving for ``/v1/realtime?duplex=1``.
|
|
|
|
Everything about a duplex *session* lives inside the engine (see
|
|
``vllm_omni.entrypoints.duplex_omni``). This handler only does transport:
|
|
|
|
* accept the socket and run the Realtime handshake (``session.update`` opens a
|
|
session through ``DuplexOmni.open_session``; ``session.resume`` re-attaches
|
|
to an existing one);
|
|
* translate each client JSON event into a ``DuplexCommand`` and submit it in
|
|
arrival order;
|
|
* pump the session's typed events to the wire, journaling them for resume;
|
|
* manage attachments: resume tokens, replay, takeover, detach on disconnect.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from collections.abc import Mapping
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass, replace
|
|
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.engine.duplex.commands import DuplexCommand, DuplexCommandError
|
|
from vllm_omni.engine.duplex.events import (
|
|
DuplexEvent,
|
|
SessionClosed,
|
|
SessionCreated,
|
|
SessionReplaced,
|
|
SessionResumed,
|
|
SessionResyncRequired,
|
|
)
|
|
from vllm_omni.engine.duplex.messages import DuplexSessionError
|
|
from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request
|
|
from vllm_omni.entrypoints.duplex.session_attachment import (
|
|
DuplexJournalGapError,
|
|
DuplexJournalOverflowError,
|
|
DuplexSessionAttachmentRegistry,
|
|
InvalidResumeTokenError,
|
|
ResumeToken,
|
|
)
|
|
from vllm_omni.entrypoints.duplex.websocket import (
|
|
MAX_EVENT_BYTES,
|
|
SendJson,
|
|
attachment_callbacks,
|
|
receive_text_with_timeout,
|
|
)
|
|
from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle
|
|
from vllm_omni.protocol.duplex import RealtimeInputDefaults
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
__all__ = ["OmniDuplexSessionHandler"]
|
|
|
|
_DEFAULT_CONFIG_TIMEOUT_S = 10.0
|
|
_DEFAULT_IDLE_TIMEOUT_S = 300.0
|
|
#: How long the endpoint waits for the session pump to deliver its terminal
|
|
#: event once the session is closed. Bounded so a wedged pump cannot pin the
|
|
#: connection open; the pump itself closes the socket on its way out.
|
|
_PUMP_DRAIN_TIMEOUT_S = 5.0
|
|
#: session.created is journaled by the registry when it hands out the resume credential.
|
|
_UNJOURNALED_EVENTS = (SessionCreated,)
|
|
|
|
|
|
@dataclass
|
|
class _Attachment:
|
|
"""What one websocket connection knows about the session it is attached to."""
|
|
|
|
handle: DuplexSessionHandle
|
|
generation: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _ResumeCredentials:
|
|
"""Transport credentials stamped onto ``session.created`` for a resumable session."""
|
|
|
|
attachment_generation: int
|
|
resume_token: str
|
|
|
|
|
|
class OmniDuplexSessionHandler:
|
|
"""WebSocket transport for engine-resident duplex sessions."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
duplex_omni: DuplexOmni,
|
|
config_timeout_s: float = _DEFAULT_CONFIG_TIMEOUT_S,
|
|
idle_timeout_s: float = _DEFAULT_IDLE_TIMEOUT_S,
|
|
) -> None:
|
|
self._omni = duplex_omni
|
|
self._config_timeout_s = config_timeout_s
|
|
self._idle_timeout_s = idle_timeout_s
|
|
runtime_config = duplex_omni.duplex_session_config
|
|
self._attachment_registry = DuplexSessionAttachmentRegistry(
|
|
replay_ttl_s=runtime_config.resume_replay_ttl_s,
|
|
replay_max_bytes_per_session=runtime_config.resume_replay_max_bytes_per_session,
|
|
)
|
|
self._resync_required_sessions: set[str] = set()
|
|
#: Wire defaults (input/output audio format and rate) per session. They
|
|
#: are negotiated on the session but live on the per-connection
|
|
#: envelope, so a reconnect has to be handed them back.
|
|
self._input_defaults: dict[str, RealtimeInputDefaults] = {}
|
|
self._pumps: dict[str, asyncio.Task[None]] = {}
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Entry point #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
async def handle_realtime_session(self, websocket: WebSocket) -> None:
|
|
await websocket.accept()
|
|
envelope = RealtimeEnvelope.from_query_params(websocket.query_params)
|
|
|
|
async def send_json(payload: Mapping[str, object]) -> None:
|
|
await websocket.send_json(dict(payload))
|
|
|
|
attachment: _Attachment | None = None
|
|
pending_command: dict[str, object] | None = None
|
|
try:
|
|
open_payload = envelope.initial_open_payload()
|
|
if open_payload is None:
|
|
first = await self._receive_first_message(websocket, envelope, send_json)
|
|
if first is None:
|
|
return
|
|
handshake = envelope.first_message(first)
|
|
if handshake.kind == "resume":
|
|
attachment = await self._resume(websocket, envelope, handshake.resume_payload, send_json)
|
|
if attachment is None:
|
|
return
|
|
else:
|
|
open_payload = handshake.session_payload
|
|
pending_command = handshake.pending_command_payload
|
|
if attachment is None:
|
|
assert open_payload is not None
|
|
attachment = await self._open(websocket, envelope, open_payload, send_json)
|
|
if attachment is None:
|
|
return
|
|
if pending_command is not None:
|
|
await self._submit_wire_event(attachment, envelope, pending_command, send_json)
|
|
self._input_defaults[attachment.handle.session_id] = envelope.defaults
|
|
await self._read_loop(websocket, envelope, attachment, send_json)
|
|
await self._drain_terminal_pump(attachment)
|
|
except WebSocketDisconnect:
|
|
if attachment is not None:
|
|
await self._on_disconnect(attachment)
|
|
except Exception as exc:
|
|
logger.exception("Duplex websocket session failed: %s", exc)
|
|
if attachment is not None and await self._attachment_registry.is_current_attachment(
|
|
attachment.handle.session_id, attachment.generation
|
|
):
|
|
with suppress(Exception):
|
|
await self._attachment_registry.close(attachment.handle.session_id)
|
|
with suppress(Exception):
|
|
await attachment.handle.close(reason="transport_error")
|
|
with suppress(Exception):
|
|
await send_json(envelope.error_payload("realtime_input_failed", str(exc)))
|
|
|
|
async def _receive_first_message(
|
|
self, websocket: WebSocket, envelope: RealtimeEnvelope, send_json: SendJson
|
|
) -> dict[str, object] | None:
|
|
raw = await receive_text_with_timeout(websocket, self._config_timeout_s)
|
|
if raw is None:
|
|
await send_json(envelope.error_payload("config_timeout", "Timeout waiting for session.update"))
|
|
return None
|
|
try:
|
|
first = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
await send_json(envelope.error_payload("invalid_json", "Invalid JSON in session.update"))
|
|
return None
|
|
if not isinstance(first, dict):
|
|
await send_json(envelope.error_payload("bad_event", "Realtime event must be a JSON object"))
|
|
return None
|
|
return first
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Handshake #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
async def _open(
|
|
self,
|
|
websocket: WebSocket,
|
|
envelope: RealtimeEnvelope,
|
|
session_payload: Mapping[str, object],
|
|
send_json: SendJson,
|
|
) -> _Attachment | None:
|
|
# Any ``session_id`` / ``id`` in the payload is ignored: the server
|
|
# allocates the id and announces it in ``session.created``.
|
|
try:
|
|
handle = await self._omni.open_session(session_payload)
|
|
except DuplexSessionError as exc:
|
|
await send_json(envelope.error_payload(exc.code, str(exc)))
|
|
return None
|
|
except (TypeError, ValueError) as exc:
|
|
await send_json(envelope.error_payload("invalid_session_config", str(exc)))
|
|
return None
|
|
attachment_send, attachment_close = attachment_callbacks(websocket)
|
|
created = await self._attachment_registry.create(
|
|
handle.session_id,
|
|
send=attachment_send,
|
|
close=attachment_close,
|
|
)
|
|
resume_supported = bool(handle.capabilities.supports_session_resume)
|
|
attachment = _Attachment(handle=handle, generation=created.attachment_generation)
|
|
credentials: _ResumeCredentials | None = None
|
|
if resume_supported:
|
|
credentials = _ResumeCredentials(
|
|
attachment_generation=created.attachment_generation,
|
|
resume_token=created.resume_token.plaintext,
|
|
)
|
|
self._start_pump(handle, credentials)
|
|
return attachment
|
|
|
|
async def _resume(
|
|
self,
|
|
websocket: WebSocket,
|
|
envelope: RealtimeEnvelope,
|
|
event: Mapping[str, object],
|
|
send_json: SendJson,
|
|
) -> _Attachment | None:
|
|
request = parse_resume_request(event)
|
|
if request is None:
|
|
await send_json(
|
|
envelope.error_payload(
|
|
"invalid_session_resume",
|
|
"session.resume requires session_id, resume_token, and a non-negative event sequence",
|
|
)
|
|
)
|
|
return None
|
|
session_id = request.session_id
|
|
handle = self._omni.get_session(session_id)
|
|
if handle is None or handle.closed:
|
|
await send_json(
|
|
envelope.error_payload("session_resume_expired", f"Unknown or expired duplex session: {session_id}")
|
|
)
|
|
return None
|
|
if not handle.capabilities.supports_session_resume:
|
|
await send_json(
|
|
envelope.error_payload("unsupported_session_resume", f"Session does not support resume: {session_id}")
|
|
)
|
|
return None
|
|
try:
|
|
await self._attachment_registry.authenticate_resume(
|
|
session_id,
|
|
resume_token=request.resume_token,
|
|
last_received_server_event_seq=request.last_received_server_event_seq,
|
|
)
|
|
except InvalidResumeTokenError:
|
|
await send_json(envelope.error_payload("invalid_resume_token", "Invalid duplex session resume token"))
|
|
return None
|
|
except DuplexJournalGapError:
|
|
await send_json(SessionResyncRequired(session_id=session_id, reason="journal_gap").to_realtime())
|
|
return None
|
|
except (KeyError, ValueError) as exc:
|
|
await send_json(envelope.error_payload("session_resume_conflict", str(exc)))
|
|
return None
|
|
try:
|
|
await self._omni.resume_session(
|
|
session_id,
|
|
expected_lease_generation=handle.lease_generation,
|
|
)
|
|
except DuplexSessionError as exc:
|
|
await send_json(envelope.error_payload("runtime_resume_failed", str(exc)))
|
|
return None
|
|
|
|
attachment_send, attachment_close = attachment_callbacks(websocket)
|
|
|
|
def activation_payload_factory(token: ResumeToken, generation: int) -> dict[str, object]:
|
|
return SessionResumed(
|
|
session_id=session_id,
|
|
session=dict(handle.public_session),
|
|
attachment_generation=generation,
|
|
resume_token=token.plaintext,
|
|
).to_realtime()
|
|
|
|
try:
|
|
resumed = await self._attachment_registry.resume(
|
|
session_id,
|
|
resume_token=request.resume_token,
|
|
last_received_server_event_seq=request.last_received_server_event_seq,
|
|
send=attachment_send,
|
|
close=attachment_close,
|
|
activation_payload_factory=activation_payload_factory,
|
|
)
|
|
except Exception as exc:
|
|
# The engine lease was already resumed above, so a failed activation
|
|
# would leave the session attached to nothing -- unless another
|
|
# connection won the race and is attached right now. Detaching then
|
|
# would start the disconnect grace for the *winner*, which ordinary
|
|
# heartbeats do not clear. Roll back only what this attempt owns.
|
|
if not await self._attachment_registry.has_attachment(session_id):
|
|
with suppress(DuplexSessionError):
|
|
await self._omni.detach_session(session_id)
|
|
await send_json(envelope.error_payload("session_resume_conflict", str(exc)))
|
|
return None
|
|
replaced = resumed.replaced_attachment
|
|
if replaced is not None:
|
|
with suppress(Exception):
|
|
await replaced.send(
|
|
SessionReplaced(session_id=session_id, attachment_generation=replaced.generation).to_realtime()
|
|
)
|
|
with suppress(Exception):
|
|
await replaced.close("session_replaced")
|
|
# A reconnect brings a fresh envelope carrying pcm16/16 kHz wire
|
|
# defaults. The negotiated input format is a wire default, not part of
|
|
# the public session object, so it has to be carried over explicitly:
|
|
# otherwise the first append that omits format/rate is decoded as pcm16.
|
|
remembered = self._input_defaults.get(session_id)
|
|
if remembered is not None:
|
|
envelope.defaults = remembered
|
|
self._start_pump(handle, None)
|
|
return _Attachment(handle=handle, generation=resumed.attachment_generation)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Outbound pump (session-scoped, survives reconnects) #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
def _start_pump(self, handle: DuplexSessionHandle, credentials: _ResumeCredentials | None) -> None:
|
|
existing = self._pumps.get(handle.session_id)
|
|
if existing is not None and not existing.done():
|
|
return
|
|
self._pumps[handle.session_id] = asyncio.create_task(
|
|
self._pump_events(handle, credentials),
|
|
name=f"duplex-session-pump-{handle.session_id}",
|
|
)
|
|
|
|
async def _pump_events(self, handle: DuplexSessionHandle, credentials: _ResumeCredentials | None) -> None:
|
|
session_id = handle.session_id
|
|
close_reason = "session_closed"
|
|
try:
|
|
async for event in handle.events():
|
|
if isinstance(event, SessionCreated) and credentials is not None:
|
|
event = replace(
|
|
event,
|
|
attachment_generation=credentials.attachment_generation,
|
|
resume_token=credentials.resume_token,
|
|
)
|
|
await self._send_event(session_id, event)
|
|
if isinstance(event, SessionClosed):
|
|
close_reason = event.reason or event.type
|
|
break
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.exception("Duplex event pump failed for %s: %s", session_id, exc)
|
|
finally:
|
|
self._pumps.pop(session_id, None)
|
|
self._resync_required_sessions.discard(session_id)
|
|
self._input_defaults.pop(session_id, None)
|
|
attachment = None
|
|
with suppress(Exception):
|
|
attachment = await self._attachment_registry.close(session_id)
|
|
if attachment is not None:
|
|
with suppress(Exception):
|
|
await attachment.close(close_reason)
|
|
|
|
async def _drain_terminal_pump(self, attachment: _Attachment) -> None:
|
|
"""Let the pump deliver ``session.closed`` before the endpoint returns.
|
|
|
|
``DuplexSessionHandle._deliver`` queues ``SessionClosed`` on the outbox
|
|
and marks the handle closed in the same synchronous step, so
|
|
``handle.closed`` is already true while that event is still sitting in
|
|
the queue. ``_read_loop`` stops on exactly that flag: if the engine
|
|
closes the session between two reads, the loop returns, this endpoint
|
|
returns, and the ASGI server tears the socket down with the terminal
|
|
event unsent -- the client sees an abrupt close instead of
|
|
``session.closed``. Waiting for the pump keeps the two in order; it is
|
|
the pump that sends the event and then closes with code 1000.
|
|
|
|
Only for a closed session. A takeover or a resumable disconnect leaves
|
|
the pump running for the next attachment, and must not be waited on.
|
|
"""
|
|
if not attachment.handle.closed:
|
|
return
|
|
pump = self._pumps.get(attachment.handle.session_id)
|
|
if pump is None or pump.done():
|
|
return
|
|
# shield: a drain timeout must not cancel a pump that is still the
|
|
# session's only writer.
|
|
with suppress(Exception):
|
|
await asyncio.wait_for(asyncio.shield(pump), _PUMP_DRAIN_TIMEOUT_S)
|
|
|
|
async def _send_event(self, session_id: str, event: DuplexEvent) -> None:
|
|
payload = event.to_realtime()
|
|
journal = not isinstance(event, _UNJOURNALED_EVENTS) and session_id not in self._resync_required_sessions
|
|
try:
|
|
try:
|
|
await self._attachment_registry.send_event(session_id, payload, journal=journal)
|
|
except DuplexJournalOverflowError:
|
|
first_overflow = session_id not in self._resync_required_sessions
|
|
self._resync_required_sessions.add(session_id)
|
|
if first_overflow:
|
|
resync = SessionResyncRequired(session_id=session_id, reason="journal_overflow")
|
|
await self._attachment_registry.send_event(session_id, resync.to_realtime(), journal=False)
|
|
await self._attachment_registry.send_event(session_id, payload, journal=False)
|
|
except KeyError:
|
|
# Attachment already closed (takeover or teardown); the journal is gone.
|
|
pass
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
# The socket died mid-send. The session is alive in the engine:
|
|
# treat it as a disconnect of the current attachment (the pump
|
|
# keeps journaling for a resume) instead of tearing the session down.
|
|
logger.info("Duplex transport send failed for %s: %s", session_id, exc)
|
|
await self._detach_current_attachment(session_id)
|
|
|
|
async def _detach_current_attachment(self, session_id: str) -> None:
|
|
"""Disconnect semantics for the socket currently attached to ``session_id``.
|
|
|
|
The pump is session-scoped and outlives any one connection, so it can
|
|
only name the session: the registry detaches whichever socket is
|
|
attached, which is the one whose send just failed.
|
|
"""
|
|
handle = self._omni.get_session(session_id)
|
|
if handle is None or handle.closed:
|
|
return
|
|
if handle.capabilities.supports_session_resume:
|
|
if await self._attachment_registry.detach(session_id, attachment_generation=None):
|
|
with suppress(DuplexSessionError):
|
|
await self._omni.detach_session(session_id)
|
|
return
|
|
with suppress(DuplexSessionError):
|
|
await handle.close(reason="disconnect")
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Inbound #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
async def _read_loop(
|
|
self,
|
|
websocket: WebSocket,
|
|
envelope: RealtimeEnvelope,
|
|
attachment: _Attachment,
|
|
send_json: SendJson,
|
|
) -> None:
|
|
handle = attachment.handle
|
|
idle_timeout_s = self._idle_timeout_s
|
|
config_timeout = handle.public_session.get("idle_timeout_s") if handle.public_session else None
|
|
if isinstance(config_timeout, int | float) and config_timeout > 0:
|
|
idle_timeout_s = float(config_timeout)
|
|
while not handle.closed:
|
|
raw = await receive_text_with_timeout(websocket, idle_timeout_s)
|
|
if not await self._attachment_registry.is_current_attachment(handle.session_id, attachment.generation):
|
|
# A newer connection took the session over; this socket is done.
|
|
return
|
|
if raw is None:
|
|
# Nothing from this socket for the idle window: same as a
|
|
# disconnect (engine grace for resumable sessions); the engine
|
|
# lease decides whether the session itself expires.
|
|
await self._on_disconnect(attachment)
|
|
return
|
|
if len(raw.encode("utf-8")) > MAX_EVENT_BYTES:
|
|
await send_json(envelope.error_payload("event_too_large", "Duplex event too large"))
|
|
continue
|
|
try:
|
|
payload = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
await send_json(envelope.error_payload("invalid_json", "Invalid JSON event"))
|
|
continue
|
|
if not isinstance(payload, dict) or not isinstance(payload.get("type"), str):
|
|
await send_json(envelope.error_payload("bad_event", "Duplex event must be a JSON object with a type"))
|
|
continue
|
|
if envelope.is_envelope_event(payload):
|
|
await self._handle_envelope_event(attachment, envelope, payload, send_json)
|
|
continue
|
|
await self._submit_wire_event(attachment, envelope, payload, send_json)
|
|
|
|
async def _handle_envelope_event(
|
|
self,
|
|
attachment: _Attachment,
|
|
envelope: RealtimeEnvelope,
|
|
payload: dict[str, object],
|
|
send_json: SendJson,
|
|
) -> None:
|
|
handle = attachment.handle
|
|
event_type = payload.get("type")
|
|
event_id = payload.get("event_id")
|
|
if event_type == "session.event_ack":
|
|
acknowledged = payload.get("server_event_seq")
|
|
if not isinstance(acknowledged, int) or acknowledged < 0:
|
|
message = "session.event_ack requires a non-negative server_event_seq"
|
|
await send_json(envelope.error_payload("invalid_event_ack", message, event_id=event_id))
|
|
return
|
|
try:
|
|
await self._attachment_registry.acknowledge(handle.session_id, acknowledged)
|
|
except ValueError as exc:
|
|
await send_json(envelope.error_payload("invalid_event_ack", str(exc), event_id=event_id))
|
|
return
|
|
message = "session.resume must be the first message of a new connection"
|
|
await send_json(envelope.error_payload("unsupported_session_resume", message, event_id=event_id))
|
|
|
|
async def _submit_wire_event(
|
|
self,
|
|
attachment: _Attachment,
|
|
envelope: RealtimeEnvelope,
|
|
payload: dict[str, object],
|
|
send_json: SendJson,
|
|
) -> None:
|
|
try:
|
|
command = envelope.translate(payload)
|
|
except DuplexCommandError as exc:
|
|
await send_json(envelope.command_error_payload(exc))
|
|
return
|
|
# ``translate`` folds a session.update's audio settings into the
|
|
# envelope; remember them so a later reconnect starts where this left off.
|
|
self._input_defaults[attachment.handle.session_id] = envelope.defaults
|
|
await self._submit_command(attachment, envelope, command, send_json)
|
|
|
|
async def _submit_command(
|
|
self,
|
|
attachment: _Attachment,
|
|
envelope: RealtimeEnvelope,
|
|
command: DuplexCommand,
|
|
send_json: SendJson,
|
|
) -> None:
|
|
try:
|
|
await attachment.handle.submit(command)
|
|
except DuplexSessionError as exc:
|
|
await send_json(envelope.error_payload(exc.code, str(exc), event_id=command.event_id))
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Disconnect #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
async def _on_disconnect(self, attachment: _Attachment) -> None:
|
|
handle = attachment.handle
|
|
if handle.closed:
|
|
return
|
|
if handle.capabilities.supports_session_resume:
|
|
if await self._attachment_registry.detach(handle.session_id, attachment_generation=attachment.generation):
|
|
# Engine-owned disconnect grace: expiry arrives as session.expired.
|
|
with suppress(DuplexSessionError):
|
|
await self._omni.detach_session(handle.session_id)
|
|
# Otherwise a newer connection already took the session over; the
|
|
# replaced socket's disconnect must not touch it.
|
|
return
|
|
with suppress(DuplexSessionError):
|
|
await handle.close(reason="disconnect")
|