hermes-chatto-plugin/vendor/common/chattolib/_sync/realtime.py

"""Synchronous realtime WebSocket client.

Hand-written sync twin of :mod:`chattolib.realtime` (the async module remains
the reference). It reuses that module's shared dataclasses, helpers and
exceptions verbatim and only reimplements the WebSocket I/O on top of
``websockets.sync``. Usage::

    with SyncChattoClient.login(...) as client:
        for event in stream_events(client):
            print(event.kind, event.payload)

Requires the ``chattolib[realtime]`` extra (``websockets``).
"""

from __future__ import annotations

from collections.abc import Iterator
from typing import TYPE_CHECKING, Any

from chattolib import _pb  # noqa: F401 — installs pb import path
from chattolib.exceptions import ChattoError
from chattolib.realtime import (
    REALTIME_PROTOCOL_VERSION,
    ChattoRealtimeCloseError,
    ChattoRealtimeError,
    RealtimeEvent,
    ServerHello,
    _raise_for_control,
    _wrap_event,
    realtime_url,
)

if TYPE_CHECKING:
    from chattolib._sync.client import SyncChattoClient


class SyncRealtimeConnection:
    """Synchronous twin of :class:`chattolib.realtime.RealtimeConnection`."""

    def __init__(
        self,
        client: SyncChattoClient,
        *,
        protocol_version: int = REALTIME_PROTOCOL_VERSION,
    ) -> None:
        self._client = client
        self._protocol_version = protocol_version
        self._ws: Any = None
        self._server_hello: ServerHello | None = None

    @property
    def server_hello(self) -> ServerHello | None:
        return self._server_hello

    def __enter__(self) -> SyncRealtimeConnection:
        self.connect()
        return self

    def __exit__(self, *exc: Any) -> None:
        self.close()

    def connect(self) -> None:
        try:
            from websockets.sync.client import connect as ws_connect
        except ImportError as exc:  # pragma: no cover
            raise ChattoError(
                "The realtime channel requires the 'websockets' package. "
                "Install with `pip install chattolib[realtime]`."
            ) from exc

        from chattolib._pb.chatto.realtime.v1 import realtime_pb2

        url = realtime_url(self._client.base_url)
        headers: dict[str, str] = {}
        if self._client.session_cookie:
            headers["Cookie"] = f"chatto_session={self._client.session_cookie}"

        self._ws = ws_connect(url, additional_headers=headers)

        client_hello = realtime_pb2.RealtimeClientFrame()
        client_hello.hello.protocol_version = self._protocol_version
        if self._client.token:
            client_hello.hello.bearer_token = self._client.token
        self._ws.send(client_hello.SerializeToString())

        first = self._ws.recv()
        first_frame = realtime_pb2.RealtimeServerFrame()
        first_frame.ParseFromString(first)
        if first_frame.WhichOneof("frame") != "hello":
            _raise_for_control(first_frame)
            raise ChattoRealtimeError(
                "unexpected_frame",
                f"expected server hello, got {first_frame.WhichOneof('frame')!r}",
            )
        hello = first_frame.hello
        self._server_hello = ServerHello(
            protocol_version=hello.protocol_version,
            server_version=hello.server_version,
            heartbeat_interval_seconds=hello.heartbeat_interval_seconds,
            capabilities=list(hello.capabilities),
        )

        subscribe = realtime_pb2.RealtimeClientFrame()
        subscribe.subscribe_events.SetInParent()
        self._ws.send(subscribe.SerializeToString())

    def close(self) -> None:
        if self._ws is None:
            return
        try:
            self._ws.close()
        finally:
            self._ws = None

    def ping(self, nonce: str = "") -> None:
        """Send a client ping. The server replies with a matching pong."""
        if self._ws is None:
            raise ChattoError("realtime connection is closed")
        from chattolib._pb.chatto.realtime.v1 import realtime_pb2

        frame = realtime_pb2.RealtimeClientFrame()
        frame.ping.nonce = nonce
        self._ws.send(frame.SerializeToString())

    def events(self) -> Iterator[RealtimeEvent]:
        """Yield decoded live events until the connection closes."""
        if self._ws is None:
            raise ChattoError("realtime connection is not open")
        from chattolib._pb.chatto.realtime.v1 import realtime_pb2

        for raw in self._ws:
            if isinstance(raw, str):
                raise ChattoRealtimeError(
                    "unexpected_text_frame",
                    "server sent a text frame; realtime protocol expects binary",
                )
            frame = realtime_pb2.RealtimeServerFrame()
            frame.ParseFromString(raw)
            case = frame.WhichOneof("frame")
            if case == "event":
                yield _wrap_event(frame.event)
            elif case in ("heartbeat", "pong", "subscribed"):
                continue
            elif case == "error":
                err = frame.error
                exc = ChattoRealtimeError(err.code, err.message, fatal=err.fatal)
                if err.fatal:
                    raise exc
                continue
            elif case == "close":
                close = frame.close
                raise ChattoRealtimeCloseError(
                    close.code,
                    close.message,
                    reconnect=close.reconnect,
                    retry_after_ms=close.retry_after_ms,
                )
            else:
                raise ChattoRealtimeError(
                    "unexpected_frame",
                    f"unknown server frame: {case!r}",
                )


def stream_events(
    client: SyncChattoClient,
    *,
    protocol_version: int = REALTIME_PROTOCOL_VERSION,
) -> Iterator[RealtimeEvent]:
    """Open a realtime connection and yield events until the server closes."""
    conn = SyncRealtimeConnection(client, protocol_version=protocol_version)
    try:
        conn.connect()
        yield from conn.events()
    finally:
        conn.close()


__all__ = ["SyncRealtimeConnection", "stream_events"]