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

# mypy: disable-error-code="no-any-return"
"""Shared plumbing for the ConnectRPC-based transport.

Chatto's public API is a ConnectRPC service surface. ``chattolib.client``
speaks to it through generated service stubs (see ``chattolib._pb``) driven
by the official ``connectrpc`` Python package.

This module exposes:

* :func:`build_service_clients` — one call, one ``base_url`` argument,
  returns a ``ServiceClients`` object with a typed field per Chatto service.
* :func:`translate_connect_error` — translates a
  ``connectrpc.errors.ConnectError`` into the library's public exception
  hierarchy (:class:`chattolib.exceptions.ChattoAuthError` /
  :class:`chattolib.exceptions.ChattoConnectError`).
* :func:`pb_to_dict` — turns a protobuf response into the camelCase JSON
  shape that the existing ``types.py`` dataclass parsers already accept.

Keeping the parsers dict-driven means the migration from Connect-JSON to
Connect-binary transport doesn't ripple through the entire public API.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from connectrpc.code import Code
from connectrpc.compat import google_protobuf_binary_codec
from connectrpc.errors import ConnectError
from google.protobuf.json_format import MessageToDict
from google.protobuf.message import Message

# Sync client twins (generated alongside the async ones in every *_connect.py).
from chattolib._pb.chatto.admin.v1.diagnostics_connect import (
    AdminDiagnosticsServiceClient,
    AdminDiagnosticsServiceClientSync,
)
from chattolib._pb.chatto.admin.v1.event_log_connect import (
    AdminEventLogServiceClient,
    AdminEventLogServiceClientSync,
)
from chattolib._pb.chatto.admin.v1.members_connect import (
    AdminUserServiceClient,
    AdminUserServiceClientSync,
)
from chattolib._pb.chatto.admin.v1.permissions_connect import (
    AdminPermissionServiceClient,
    AdminPermissionServiceClientSync,
)
from chattolib._pb.chatto.admin.v1.roles_connect import (
    AdminRoleServiceClient,
    AdminRoleServiceClientSync,
)
from chattolib._pb.chatto.admin.v1.room_layout_connect import (
    AdminRoomLayoutServiceClient,
    AdminRoomLayoutServiceClientSync,
)
from chattolib._pb.chatto.admin.v1.server_connect import (
    AdminServerServiceClient,
    AdminServerServiceClientSync,
)
from chattolib._pb.chatto.api.v1.account_connect import (
    MyAccountServiceClient,
    MyAccountServiceClientSync,
)
from chattolib._pb.chatto.api.v1.asset_uploads_connect import (
    AssetUploadServiceClient,
    AssetUploadServiceClientSync,
)
from chattolib._pb.chatto.api.v1.attachments_connect import (
    AssetServiceClient,
    AssetServiceClientSync,
)
from chattolib._pb.chatto.api.v1.member_directory_connect import (
    UserServiceClient,
    UserServiceClientSync,
)
from chattolib._pb.chatto.api.v1.messages_connect import (
    MessageServiceClient,
    MessageServiceClientSync,
)
from chattolib._pb.chatto.api.v1.notification_preferences_connect import (
    NotificationPreferencesServiceClient,
    NotificationPreferencesServiceClientSync,
)
from chattolib._pb.chatto.api.v1.notifications_connect import (
    NotificationServiceClient,
    NotificationServiceClientSync,
)
from chattolib._pb.chatto.api.v1.push_notifications_connect import (
    PushNotificationServiceClient,
    PushNotificationServiceClientSync,
)
from chattolib._pb.chatto.api.v1.roles_connect import RoleServiceClient, RoleServiceClientSync
from chattolib._pb.chatto.api.v1.room_directory_connect import (
    RoomDirectoryServiceClient,
    RoomDirectoryServiceClientSync,
)
from chattolib._pb.chatto.api.v1.rooms_connect import RoomServiceClient, RoomServiceClientSync
from chattolib._pb.chatto.api.v1.server_state_connect import (
    ServerServiceClient,
    ServerServiceClientSync,
)
from chattolib._pb.chatto.api.v1.threads_connect import ThreadServiceClient, ThreadServiceClientSync
from chattolib._pb.chatto.api.v1.viewer_connect import ViewerServiceClient, ViewerServiceClientSync
from chattolib._pb.chatto.api.v1.voice_calls_connect import (
    VoiceCallServiceClient,
    VoiceCallServiceClientSync,
)
from chattolib._pb.chatto.auth.v1.external_identity_auth_connect import (
    ExternalIdentityAuthServiceClient,
    ExternalIdentityAuthServiceClientSync,
)
from chattolib._pb.chatto.discovery.v1.server_connect import (
    ServerDiscoveryServiceClient,
    ServerDiscoveryServiceClientSync,
)
from chattolib.exceptions import ChattoAuthError, ChattoConnectError

CONNECT_PREFIX = "/api/connect"


@dataclass
class ServiceClients:
    """Typed bundle of ConnectRPC service clients used by ``ChattoClient``."""

    server_discovery: ServerDiscoveryServiceClient
    server: ServerServiceClient
    viewer: ViewerServiceClient
    account: MyAccountServiceClient
    users: UserServiceClient
    roles: RoleServiceClient
    room_directory: RoomDirectoryServiceClient
    rooms: RoomServiceClient
    messages: MessageServiceClient
    threads: ThreadServiceClient
    notifications: NotificationServiceClient
    notification_prefs: NotificationPreferencesServiceClient
    push: PushNotificationServiceClient
    assets: AssetServiceClient
    asset_uploads: AssetUploadServiceClient
    voice_calls: VoiceCallServiceClient
    external_auth: ExternalIdentityAuthServiceClient
    admin_server: AdminServerServiceClient
    admin_room_layout: AdminRoomLayoutServiceClient
    admin_users: AdminUserServiceClient
    admin_roles: AdminRoleServiceClient
    admin_event_log: AdminEventLogServiceClient
    admin_diagnostics: AdminDiagnosticsServiceClient
    admin_permissions: AdminPermissionServiceClient

    async def close(self) -> None:
        for name in self.__dataclass_fields__:
            client = getattr(self, name)
            await client.close()


def build_service_clients(base_url: str) -> ServiceClients:
    """Instantiate one service client per Chatto Connect service.

    ``base_url`` is the server root (e.g. ``https://chat.chatto.run``); the
    ConnectRPC prefix is appended by this function.
    """
    address = f"{base_url.rstrip('/')}{CONNECT_PREFIX}"
    codec = google_protobuf_binary_codec()

    def make(cls: Any) -> Any:
        return cls(address, codec=codec)

    return ServiceClients(
        server_discovery=make(ServerDiscoveryServiceClient),
        server=make(ServerServiceClient),
        viewer=make(ViewerServiceClient),
        account=make(MyAccountServiceClient),
        users=make(UserServiceClient),
        roles=make(RoleServiceClient),
        room_directory=make(RoomDirectoryServiceClient),
        rooms=make(RoomServiceClient),
        messages=make(MessageServiceClient),
        threads=make(ThreadServiceClient),
        notifications=make(NotificationServiceClient),
        notification_prefs=make(NotificationPreferencesServiceClient),
        push=make(PushNotificationServiceClient),
        assets=make(AssetServiceClient),
        asset_uploads=make(AssetUploadServiceClient),
        voice_calls=make(VoiceCallServiceClient),
        external_auth=make(ExternalIdentityAuthServiceClient),
        admin_server=make(AdminServerServiceClient),
        admin_room_layout=make(AdminRoomLayoutServiceClient),
        admin_users=make(AdminUserServiceClient),
        admin_roles=make(AdminRoleServiceClient),
        admin_event_log=make(AdminEventLogServiceClient),
        admin_diagnostics=make(AdminDiagnosticsServiceClient),
        admin_permissions=make(AdminPermissionServiceClient),
    )


def translate_connect_error(exc: ConnectError) -> Exception:
    """Convert a ``connectrpc`` error into chattolib's exception hierarchy."""
    if exc.code == Code.UNAUTHENTICATED:
        return ChattoAuthError(str(exc))
    return ChattoConnectError(
        code=exc.code.name.lower(),
        message=str(exc),
    )


# --- Synchronous transport -------------------------------------------------
#
# The async client translates ConnectError via ``ChattoClient._rpc``, which
# awaits the coroutine inside a try/except. The sync stubs raise at call time
# (before the generated ``_rpc`` passthrough runs), so error translation is
# done here at the transport via a ConnectRPC unary interceptor instead.


class _ErrorTranslatingInterceptorSync:
    """Structural ``connectrpc.UnaryInterceptorSync`` that maps ConnectError."""

    def intercept_unary_sync(self, call_next: Any, request: Any, ctx: Any) -> Any:
        try:
            return call_next(request, ctx)
        except ConnectError as exc:
            raise translate_connect_error(exc) from exc


@dataclass
class ServiceClientsSync:
    """Synchronous twin of :class:`ServiceClients`."""

    server_discovery: ServerDiscoveryServiceClientSync
    server: ServerServiceClientSync
    viewer: ViewerServiceClientSync
    account: MyAccountServiceClientSync
    users: UserServiceClientSync
    roles: RoleServiceClientSync
    room_directory: RoomDirectoryServiceClientSync
    rooms: RoomServiceClientSync
    messages: MessageServiceClientSync
    threads: ThreadServiceClientSync
    notifications: NotificationServiceClientSync
    notification_prefs: NotificationPreferencesServiceClientSync
    push: PushNotificationServiceClientSync
    assets: AssetServiceClientSync
    asset_uploads: AssetUploadServiceClientSync
    voice_calls: VoiceCallServiceClientSync
    external_auth: ExternalIdentityAuthServiceClientSync
    admin_server: AdminServerServiceClientSync
    admin_room_layout: AdminRoomLayoutServiceClientSync
    admin_users: AdminUserServiceClientSync
    admin_roles: AdminRoleServiceClientSync
    admin_event_log: AdminEventLogServiceClientSync
    admin_diagnostics: AdminDiagnosticsServiceClientSync
    admin_permissions: AdminPermissionServiceClientSync

    def close(self) -> None:
        for name in self.__dataclass_fields__:
            getattr(self, name).close()


def build_service_clients_sync(base_url: str) -> ServiceClientsSync:
    """Synchronous twin of :func:`build_service_clients`.

    Each client carries an error-translating interceptor so ConnectError
    surfaces as ``ChattoConnectError`` / ``ChattoAuthError`` exactly as the
    async path does.
    """
    address = f"{base_url.rstrip('/')}{CONNECT_PREFIX}"
    codec = google_protobuf_binary_codec()
    interceptor = _ErrorTranslatingInterceptorSync()

    def make(cls: Any) -> Any:
        return cls(address, codec=codec, interceptors=(interceptor,))

    return ServiceClientsSync(
        server_discovery=make(ServerDiscoveryServiceClientSync),
        server=make(ServerServiceClientSync),
        viewer=make(ViewerServiceClientSync),
        account=make(MyAccountServiceClientSync),
        users=make(UserServiceClientSync),
        roles=make(RoleServiceClientSync),
        room_directory=make(RoomDirectoryServiceClientSync),
        rooms=make(RoomServiceClientSync),
        messages=make(MessageServiceClientSync),
        threads=make(ThreadServiceClientSync),
        notifications=make(NotificationServiceClientSync),
        notification_prefs=make(NotificationPreferencesServiceClientSync),
        push=make(PushNotificationServiceClientSync),
        assets=make(AssetServiceClientSync),
        asset_uploads=make(AssetUploadServiceClientSync),
        voice_calls=make(VoiceCallServiceClientSync),
        external_auth=make(ExternalIdentityAuthServiceClientSync),
        admin_server=make(AdminServerServiceClientSync),
        admin_room_layout=make(AdminRoomLayoutServiceClientSync),
        admin_users=make(AdminUserServiceClientSync),
        admin_roles=make(AdminRoleServiceClientSync),
        admin_event_log=make(AdminEventLogServiceClientSync),
        admin_diagnostics=make(AdminDiagnosticsServiceClientSync),
        admin_permissions=make(AdminPermissionServiceClientSync),
    )


def pb_to_dict(message: Message | None) -> dict[str, Any]:
    """Convert a protobuf message to the camelCase dict shape the parsers accept.

    ``preserving_proto_field_name=False`` gives us JSON-mapping camelCase
    keys (e.g. ``created_at`` → ``createdAt``), matching what the
    ``types.py`` dataclass parsers already consume.
    """
    if message is None:
        return {}
    return MessageToDict(
        message,
        preserving_proto_field_name=False,
        use_integers_for_enums=False,
    )