hermes-chatto-plugin/vendor/common/chattolib/bot.py

"""A small, opinionated framework for building bots on Chatto.

Chatto 0.5.0 introduced *bot accounts*: user identities flagged as bots that
authenticate with a **key** (e.g. ``cht_BK_...``) rather than a
username/password. The key is used **directly as a bearer token** — there is
no ``/auth/login`` round-trip — and it resolves to a ``User`` with a
capability grant set.

This module turns :class:`chattolib.client.ChattoClient` into the standard
library for bots by adding three things on top of the raw client:

* **Key-based login** — :meth:`Bot.login` builds an authenticated client from
  a bot key in one call.
* **An event dispatcher** — :meth:`Bot.run` opens the realtime stream and
  routes incoming events (messages, reactions, mentions, presence, typing,
  room changes) to the async handlers you register.
* **Flattened verbs** — :meth:`Bot.say`, :meth:`Bot.reply`,
  :meth:`Bot.react`, :method:`Bot.set_status`, :meth:`Bot.join_room`,
  :meth:`Bot.create_room`, and friends, so a bot's "brain" reads like
  natural language instead of protobuf plumbing.

In the protocol-4 realtime channel, live events are thin, caller-scoped hints
(identifiers rather than full resources). Where a handler needs the full
resource (e.g. the body of a posted message), the bot hydrates it on demand
through the corresponding ConnectRPC.

A minimal bot::

    import asyncio
    from chattolib.bot import Bot

    async def on_message(event):
        if event.body.startswith("!hello"):
            await event.bot.reply(event, "hi!")

    async def main():
        async with await Bot.login("cht_BK_...") as bot:
            bot.on("message", on_message)
            await bot.run()

    asyncio.run(main())

Requires the ``chattolib[realtime]`` extra for the live stream.
"""

from __future__ import annotations

import asyncio
import contextlib
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any

from chattolib import _pb  # noqa: F401 — installs the generated pb import path
from chattolib._pb.chatto.api.v1 import presence_pb2
from chattolib.client import ChattoClient
from chattolib.exceptions import ChattoError
from chattolib.realtime import (
    ChattoRealtimeCloseError,
    ChattoRealtimeError,
    RealtimeConnection,
    RealtimeEvent,
    stream_events,
)
from chattolib.types import (
    Message,
    PresenceStatus,
    Room,
    RoomGroup,
    RoomWithViewerState,
    User,
)


def _presence(value: Any) -> PresenceStatus:
    """Coerce a protobuf presence value (int or enum) to a PresenceStatus."""
    name = value.name if hasattr(value, "name") else presence_pb2.PresenceStatus.Name(int(value))
    # The protobuf name is the *value* of our StrEnum (e.g. "PRESENCE_STATUS_ONLINE"),
    # so look it up by value, not by member name.
    for member in PresenceStatus:
        if member.value == name:
            return member
    return PresenceStatus.UNSPECIFIED


__all__ = [
    "Bot",
    "BotError",
    "BotEvent",
    "BotMessageEvent",
    "BotPresenceEvent",
    "BotReactionEvent",
    "BotRoomEvent",
    "BotTypingEvent",
    "BotUserEvent",
]

# Oneof ``event`` cases that map onto a :class:`BotRoomEvent`.
_ROOM_EVENT_KINDS = frozenset(
    {
        "room_created",
        "room_updated",
        "room_deleted",
        "room_archived",
        "room_unarchived",
        "room_universal_changed",
        "room_slow_mode_changed",
        "room_threading_mode_changed",
        "user_joined_room",
        "user_left_room",
    }
)
# Oneof ``event`` cases that map onto a :class:`BotUserEvent`.
_USER_EVENT_KINDS = frozenset(
    {"user_account_created", "user_profile_changed", "user_account_deleted"}
)


class BotError(ChattoError):
    """Raised for bot-framework-level errors (bad handlers, bad keys, ...)."""


# ---------------------------------------------------------------------------
# Event types
# ---------------------------------------------------------------------------


@dataclass
class BotEvent:
    """Base class for every event the dispatcher delivers to a handler.

    ``bot`` is the owning :class:`Bot`, so a handler can act on the event
    (e.g. ``await event.bot.reply(event, "...")``) without closing over
    globals.
    """

    bot: Bot
    kind: str

    async def _noop(self) -> None:  # pragma: no cover - interface
        ...


@dataclass
class BotMessageEvent(BotEvent):
    """A new message in a room the bot can see.

    ``message`` is the full :class:`Message`, hydrated from the server when
    this event was dispatched. ``actor`` is the author's :class:`User` when
    it could be resolved; otherwise ``None``.
    """

    message: Message
    actor: User | None = None

    @property
    def room_id(self) -> str:
        return self.message.room_id

    @property
    def body(self) -> str | None:
        return self.message.body

    @property
    def is_mention(self) -> bool:
        """True when this message mentions the bot (best-effort)."""
        me = self.bot.user
        if me is None or not self.message.body:
            return False
        return f"@{me.login}" in self.message.body or me.display_name in self.message.body


@dataclass
class BotReactionEvent(BotEvent):
    """A reaction was added to (or removed from) a message."""

    room_id: str
    message_event_id: str
    emoji: str
    user_id: str
    added: bool = True


@dataclass
class BotPresenceEvent(BotEvent):
    """A user's presence status changed."""

    user_id: str
    status: PresenceStatus


@dataclass
class BotTypingEvent(BotEvent):
    """A user started (or stopped) typing in a room/thread."""

    room_id: str
    thread_root_event_id: str | None = None


@dataclass
class BotRoomEvent(BotEvent):
    """A room lifecycle change (created, updated, archived, member join/...)."""

    room: Room | None = None
    detail: str = ""  # e.g. "created", "updated", "archived", "user_joined_room"


@dataclass
class BotUserEvent(BotEvent):
    """A user account was created, changed, or deleted."""

    user: User | None = None
    removed: bool = False


# ---------------------------------------------------------------------------
# The Bot
# ---------------------------------------------------------------------------


Handler = Callable[[BotEvent], Awaitable[None]]


class Bot:
    """A Chatto bot: an authenticated client plus an event dispatcher and a
    set of convenience verbs.

    Create one with :meth:`login` (from a bot key) or :meth:`from_client`
    (wrapping an already-authenticated :class:`ChattoClient`).
    """

    def __init__(
        self,
        client: ChattoClient,
        *,
        base_url: str | None = None,
    ) -> None:
        self._client = client
        self._base_url = base_url or client.base_url
        self._user: User | None = None
        self._handlers: dict[str, list[Handler]] = {}
        self._connection: RealtimeConnection | None = None
        self._running = False

    # -- construction ------------------------------------------------------

    @classmethod
    async def login(
        cls,
        key: str,
        *,
        base_url: str | None = None,
    ) -> Bot:
        """Authenticate a bot with its **key**.

        The key (e.g. ``cht_BK_...``) is used directly as the bearer token —
        no username/password. ``base_url`` defaults to the public Chatto
        server; pass it to target a self-hosted or preview deployment.
        """
        if base_url is None:
            base_url = ChattoClient.DEFAULT_BASE_URL
        client = ChattoClient(token=key, base_url=base_url)
        bot = cls(client, base_url=base_url)
        await bot._probe_identity()
        return bot

    @classmethod
    def from_client(cls, client: ChattoClient) -> Bot:
        """Wrap an already-authenticated client (e.g. a human login)."""
        return cls(client)

    async def _probe_identity(self) -> None:
        """Resolve the bot's own identity and warn (not fail) if not a bot."""
        try:
            self._user = await self._client.me()
        except ChattoError:
            self._user = None
            return
        if self._user is not None and not self._user.is_bot:
            # Not fatal — a human client can drive the same verbs — but surface
            # it so a misconfigured key is obvious.
            import warnings

            warnings.warn(
                f"Bot key authenticated as {self._user.login!r}, which is not "
                "flagged is_bot. It will still work, but this is usually a "
                "misconfigured key.",
                stacklevel=2,
            )

    # -- identity ----------------------------------------------------------

    @property
    def client(self) -> ChattoClient:
        """The underlying :class:`ChattoClient` for any RPC not wrapped here."""
        return self._client

    @property
    def base_url(self) -> str:
        return self._base_url

    @property
    def user(self) -> User | None:
        """The bot's own :class:`User` profile, once known."""
        return self._user

    @property
    def login_name(self) -> str | None:
        return self._user.login if self._user else None

    @property
    def display_name(self) -> str | None:
        return self._user.display_name if self._user else None

    # -- lifecycle ---------------------------------------------------------

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

    async def __aexit__(self, *exc: Any) -> None:
        await self.close()

    async def close(self) -> None:
        if self._connection is not None:
            with contextlib.suppress(Exception):
                await self._connection.close()
            self._connection = None
        self._running = False
        await self._client.close()

    # -- event registration ------------------------------------------------

    def on(self, kind: str, handler: Handler) -> Handler:
        """Register ``handler`` for events of ``kind``.

        ``kind`` is one of: ``message``, ``reaction``, ``presence``,
        ``typing``, ``room``, ``user``, ``*`` (all events). ``handler`` is an
        ``async def`` taking a single :class:`BotEvent`. Returns the handler
        so ``on`` can be used as a decorator.
        """
        self._handlers.setdefault(kind, []).append(handler)
        return handler

    def off(self, kind: str, handler: Handler) -> None:
        with contextlib.suppress(ValueError):
            self._handlers[kind].remove(handler)

    async def _dispatch(self, event: BotEvent) -> None:
        handlers = list(self._handlers.get(event.kind, [])) + list(self._handlers.get("*", []))
        for handler in handlers:
            try:
                await handler(event)
            except Exception:  # noqa: BLE001 - one bad handler must not kill the loop
                import logging

                logging.exception("bot handler for %r raised", event.kind)

    # -- verbs: messaging --------------------------------------------------

    async def say(self, room_id: str, body: str = "", *, join_if_needed: bool = True) -> Message:
        """Post a message to a room. Returns the created :class:`Message`.

        With ``join_if_needed`` (the default), the bot joins the room first if
        it isn't already a member — so a bot that wants to be present
        everywhere can simply ``await bot.say(room_id, ...)`` without a
        separate join step. Pass ``join_if_needed=False`` to instead surface
        the server's ``permission_denied`` if the bot can't post.
        """
        try:
            return await self._client.post_message(room_id, body)
        except ChattoError as exc:
            if not join_if_needed or "not a member" not in str(exc):
                raise
            await self._client.join_room(room_id)
            return await self._client.post_message(room_id, body)

    async def reply(
        self,
        target: BotMessageEvent | Message,
        body: str,
        *,
        also_send_to_channel: bool = False,
    ) -> Message:
        """Reply to a message (or a :class:`BotMessageEvent`) in its room/thread."""
        if isinstance(target, BotMessageEvent):
            target = target.message
        return await self._client.post_message(
            target.room_id,
            body,
            thread_root_event_id=target.thread_root_event_id,
            in_reply_to=target.id,
            also_send_to_channel=also_send_to_channel,
        )

    async def react(self, room_id: str, message_event_id: str, emoji: str) -> bool:
        """Add an emoji reaction to a message."""
        return await self._client.add_reaction(room_id, message_event_id, emoji)

    async def unreact(self, room_id: str, message_event_id: str, emoji: str) -> bool:
        """Remove one of the bot's emoji reactions from a message."""
        return await self._client.remove_reaction(room_id, message_event_id, emoji)

    # -- verbs: presence & status -----------------------------------------

    async def set_presence(self, status: PresenceStatus) -> PresenceStatus:
        """Set the bot's presence (``ONLINE`` / ``AWAY`` / ``DO_NOT_DISTURB``)."""
        return await self._client.set_presence(status)

    async def set_status(self, emoji: str, text: str) -> dict[str, Any]:
        """Set the bot's custom status (e.g. a "working on X" note)."""
        return await self._client.set_custom_status(emoji, text)

    async def clear_status(self) -> dict[str, Any]:
        """Clear the bot's custom status."""
        return await self._client.delete_custom_status()

    # -- verbs: rooms ------------------------------------------------------

    async def join_room(self, room_id: str) -> Room:
        """Join a room. Returns the joined :class:`Room`."""
        return await self._client.join_room(room_id)

    async def join_room_group(self, group_id: str) -> list[str]:
        """Join **all** the rooms in a room group in one call.

        Mirrors the Chatto UI's one-click "join group" action. Returns the
        list of room IDs the bot is now a member of as a result.
        """
        return await self._client.join_room_group(group_id)

    async def list_room_groups(self) -> list[RoomGroup]:
        """List the room groups (and the rooms each contains) the bot can see."""
        return await self._client.list_room_groups()

    async def join_all_rooms(self) -> list[str]:
        """Join every room the bot can see, grouped the way the UI does.

        Joins each room group in one call (so a bot becomes a member of all
        the rooms in a group at once), then joins any ungrouped rooms
        individually. Returns the room IDs the bot is now a member of.
        """
        joined: list[str] = []
        groups = await self.list_room_groups()
        grouped_room_ids: set[str] = set()
        for group in groups:
            for rws in group.rooms:
                if rws.room is not None:
                    grouped_room_ids.add(rws.room.id)
            try:
                joined.extend(await self.join_room_group(group.id))
            except ChattoError:
                # A group the bot can't join is skipped, not fatal.
                continue
        # Rooms that are not part of any group.
        for rws in await self.list_rooms():
            if rws.room is None or rws.room.id in grouped_room_ids:
                continue
            if rws.viewer_state.is_member:
                continue
            try:
                joined.append((await self.join_room(rws.room.id)).id)
            except ChattoError:
                continue
        return joined

    async def leave_room(self, room_id: str) -> None:
        """Leave a room."""
        await self._client.leave_room(room_id)

    async def create_room(
        self,
        name: str,
        group_id: str,
        *,
        description: str = "",
        universal: bool = False,
    ) -> Room:
        """Create a room in a room group."""
        return await self._client.create_room(
            name, group_id, description=description, universal=universal
        )

    async def list_rooms(self) -> list[RoomWithViewerState]:
        """List the rooms the bot is a member of."""
        return await self._client.list_rooms()

    async def mark_read(self, room_id: str) -> None:
        """Mark a room as read (clears its unread state for the bot)."""
        await self._client.mark_room_as_read(room_id)

    # -- the run loop ------------------------------------------------------

    async def run(
        self,
        *,
        resume_cursor: str | None = None,
        until: asyncio.Event | None = None,
    ) -> None:
        """Connect the realtime stream and dispatch events until it closes.

        Reconnects automatically (resuming from the last received
        ``resume_cursor``) when the server sends a reconnectable close frame.
        When the server terminates the session (or sends a non-reconnectable
        close), the loop stops. Pass ``until`` to stop on an external signal.
        """
        self._running = True
        cursor = resume_cursor
        while self._running:
            if until is not None and until.is_set():
                break
            try:
                async for frame in stream_events(self._client, resume_cursor=cursor):
                    if until is not None and until.is_set():
                        break
                    if isinstance(frame, RealtimeEvent):
                        if frame.cursor:
                            cursor = frame.cursor
                        await self._handle_live(frame)
            except ChattoRealtimeCloseError as close:
                if not close.reconnect:
                    # Session terminated or otherwise unrecoverable: stop.
                    self._running = False
                    break
                # Reconnectable close: loop again, resuming from the cursor.
                continue
            except ChattoRealtimeError:
                raise
        self._running = False

    # -- live-event dispatch -----------------------------------------------

    async def _handle_live(self, frame: RealtimeEvent) -> None:
        kind = frame.kind
        if kind == "presence_changed":
            await self._dispatch(
                BotPresenceEvent(
                    bot=self,
                    kind="presence",
                    user_id=frame.actor_id or "",
                    status=_presence(frame.payload.status),
                )
            )
        elif kind == "user_typing":
            payload = frame.payload
            await self._dispatch(
                BotTypingEvent(
                    bot=self,
                    kind="typing",
                    room_id=payload.room_id,
                    thread_root_event_id=payload.thread_root_event_id or None,
                )
            )
        elif kind in ("reaction_added", "reaction_removed"):
            payload = frame.payload
            await self._dispatch(
                BotReactionEvent(
                    bot=self,
                    kind="reaction",
                    room_id=payload.room_id,
                    message_event_id=payload.message_event_id,
                    emoji=payload.emoji,
                    user_id=frame.actor_id or "",
                    added=kind == "reaction_added",
                )
            )
        elif kind == "message_posted":
            await self._on_message_posted(frame)
        elif kind in _ROOM_EVENT_KINDS:
            await self._on_room_event(frame)
        elif kind in _USER_EVENT_KINDS:
            await self._on_user_event(frame)

    async def _on_message_posted(self, frame: RealtimeEvent) -> None:
        # Protocol-4 events are thin hints; hydrate the full message so the
        # handler gets a usable :class:`Message` (including the body).
        payload = frame.payload
        msg: Message | None = None
        with contextlib.suppress(ChattoError):
            msg = await self._client.get_message(payload.room_id, frame.id)
        if msg is None:
            return
        await self._dispatch(BotMessageEvent(bot=self, kind="message", message=msg, actor=None))

    async def _on_room_event(self, frame: RealtimeEvent) -> None:
        payload = frame.payload
        room_id = getattr(payload, "room_id", "")
        if not room_id:
            return
        extra: dict[str, Any] = {}
        if frame.kind == "room_created":
            name = getattr(payload, "name", "")
            if name:
                extra["name"] = name
        room = Room.parse({"id": room_id, **extra})
        await self._dispatch(BotRoomEvent(bot=self, kind="room", room=room, detail=frame.kind))

    async def _on_user_event(self, frame: RealtimeEvent) -> None:
        user_id = getattr(frame.payload, "user_id", "") or frame.actor_id or ""
        removed = frame.kind == "user_account_deleted"
        await self._dispatch(
            BotUserEvent(
                bot=self,
                kind="user",
                user=User.parse({"id": user_id}) if user_id else None,
                removed=removed,
            )
        )