orchestrata_ultimate.gateway.protocol

Reference › orchestrata_ultimate › gateway › protocol

protocol holds MessageProcessor, which decodes inbound WebSocket frames and pushes outbound events. server is its importer; it is constructed with the rpc handler, the session connection manager, the socket map and an optional pubsub manager. See gateway.

process validates its arguments (TypeError or ValueError), counts the frame in metrics, and parses it into a WSMessage from models. Unparseable input gets an Invalid message reply. A HEARTBEAT frame updates the connection's heartbeat; only RPC_REQUEST frames go on to _dispatch_rpc, which calls RPCHandler.dispatch and replies with an RPC_RESPONSE carrying the result, or an ERROR message if RPCError is raised. A publish call is also fanned out to other clients through _broadcast_legacy.

broadcast is the outbound path: it asks the pub/sub manager for connections whose topic pattern and filters match, increments a per-topic sequence counter, and sends each a JSON-RPC 2.0 notification with no id whose params hold the sequence, a timestamp and the payload. It returns the number of clients reached and logs, without raising, any send that fails. push_to_client sends a notification to one connection.

At a glance

Class diagram

Dependency diagram

API

Classes

MessageProcessor

Line 27

Handles inbound messages after authentication.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/gateway/protocol.py

"""WebSocket message processing and broadcast logic.

Supports both internal WSMessage format and JSON-RPC 2.0
notification pushes with per-topic sequence numbers.
"""

import json
import logging
from collections import defaultdict
from datetime import UTC, datetime
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.gateway.exceptions import RPCError
from orchestrata_ultimate.gateway.metrics import message_received, message_sent
from orchestrata_ultimate.gateway.models import (
    MessageType,
    WSMessage,
)
from orchestrata_ultimate.gateway.pubsub import PubSubManager
from orchestrata_ultimate.gateway.rpc import RPCHandler
from orchestrata_ultimate.gateway.session import ConnectionManager

logger = logging.getLogger(__name__)


class MessageProcessor:
    """Handles inbound messages after authentication.

    Also provides JSON-RPC 2.0 notification broadcasting
    for server-push event delivery to subscribers.
    """

    def __init__(
        self,
        rpc: RPCHandler,
        connections: ConnectionManager,
        ws_map: dict[str, Any],
        pubsub: PubSubManager | None = None,
    ) -> None:
        """Wire processor to gateway internals.

        Args:
            rpc: RPC method dispatcher.
            connections: Connection lifecycle manager.
            ws_map: Map of connection_id to websocket.
            pubsub: Optional pub/sub manager for broadcasts.
        """
        self._rpc = rpc
        self._connections = connections
        self._ws_map = ws_map
        self._pubsub = pubsub
        self._sequence_counters: dict[str, int] = defaultdict(int)

    @traces(spec="014", story="US-011", reqs=["REQ-014-168"])
    async def process(self, conn_id: str, raw: str) -> None:
        """Parse and route one inbound frame.

        Args:
            conn_id: Originating connection ID.
            raw: Raw JSON string from websocket.

        Raises:
            TypeError: If conn_id or raw are not strings.
            ValueError: If conn_id is empty.
        """
        if not isinstance(conn_id, str):
            raise TypeError("conn_id must be a str")
        if not conn_id:
            raise ValueError("conn_id must not be empty")
        if not isinstance(raw, str):
            raise TypeError("raw must be a str")
        # REQ-014-144: one inbound frame received, valid or not.
        message_received()
        try:
            msg = WSMessage.from_dict(json.loads(raw))
        except (
            json.JSONDecodeError,
            KeyError,
            ValueError,
        ):
            await self._send_to(conn_id, "Invalid message")
            return
        if msg.msg_type == MessageType.HEARTBEAT:
            await self._connections.update_heartbeat(conn_id)
            return
        if msg.msg_type != MessageType.RPC_REQUEST:
            return
        await self._dispatch_rpc(conn_id, msg)

    async def _dispatch_rpc(self, conn_id: str, msg: WSMessage) -> None:
        """Dispatch an RPC request and send the response.

        Args:
            conn_id: Originating connection ID.
            msg: Parsed WSMessage with RPC payload.
        """
        method = msg.payload.get("method", "")
        params = msg.payload.get("params", {})
        try:
            result = await self._rpc.dispatch(method, conn_id, params)
            if method == "publish":
                await self._broadcast_legacy(conn_id, result)
            resp = WSMessage(
                msg_type=MessageType.RPC_RESPONSE,
                payload={"result": result},
                correlation_id=msg.msg_id,
            )
        except RPCError as exc:
            resp = WSMessage(
                msg_type=MessageType.ERROR,
                payload={"error": str(exc)},
                correlation_id=msg.msg_id,
            )
        await self._send_to(conn_id, json.dumps(resp.to_dict()))

    @traces(spec="014", story="US-008", reqs=["REQ-014-131"])
    @traces(spec="014", story="US-012", reqs=["REQ-014-193"])
    async def broadcast(
        self,
        topic: str,
        event_payload: dict[str, Any],
    ) -> int:
        """Broadcast a JSON-RPC 2.0 notification to subscribers.

        Sends a server-push notification (no id field) to all
        connections subscribed to the given topic. Includes a
        monotonically increasing sequence number per topic.

        Args:
            topic: Event topic name (e.g. "cxp.traffic").
            event_payload: Event data to include in params.

        Returns:
            Number of clients the event was delivered to.
        """
        if not self._pubsub:
            return 0
        recipients = await self._pubsub.get_matching_connections(
            topic, payload=event_payload
        )
        if not recipients:
            return 0
        self._sequence_counters[topic] += 1
        sequence = self._sequence_counters[topic]
        notification = {
            "jsonrpc": "2.0",
            "method": topic,
            "params": {
                "sequence": sequence,
                "timestamp": datetime.now(UTC).isoformat(),
                "payload": event_payload,
            },
        }
        payload_str = json.dumps(notification)
        delivered = 0
        for cid in recipients:
            ws = self._ws_map.get(cid)
            if ws:
                try:
                    await ws.send(payload_str)
                    delivered += 1
                    message_sent()
                except Exception:
                    logger.exception("Broadcast to %s failed", cid)
        return delivered

    @traces(spec="014", story="US-011", reqs=["REQ-014-174"])
    async def push_to_client(
        self,
        connection_id: str,
        method: str,
        params: dict[str, Any],
    ) -> bool:
        """Send a JSON-RPC notification to a specific client.

        Args:
            connection_id: Target client connection ID.
            method: Notification method name.
            params: Notification parameters.

        Returns:
            True if delivered, False otherwise.
        """
        ws = self._ws_map.get(connection_id)
        if not ws:
            return False
        notification = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        }
        try:
            await ws.send(json.dumps(notification))
            message_sent()
            return True
        except Exception:
            logger.exception("Push to %s failed", connection_id)
            return False

    async def _broadcast_legacy(
        self,
        sender_id: str,
        result: dict[str, Any],
    ) -> None:
        """Push an event to matching subscribers (legacy format).

        Args:
            sender_id: Connection that published.
            result: Publish result with topic, data, recipients.
        """
        topic = result.get("topic", "")
        data = result.get("data", {})
        recipients = result.get("recipients", [])
        event = WSMessage(
            msg_type=MessageType.EVENT,
            payload={"topic": topic, "data": data},
        )
        payload = json.dumps(event.to_dict())
        for cid in recipients:
            if cid == sender_id:
                continue
            ws = self._ws_map.get(cid)
            if ws:
                try:
                    await ws.send(payload)
                    message_sent()
                except Exception:
                    logger.exception("Broadcast to %s failed", cid)

    async def _send_to(self, conn_id: str, text: str) -> None:
        """Send text to a specific connection.

        Args:
            conn_id: Target connection ID.
            text: Raw text to send.
        """
        ws = self._ws_map.get(conn_id)
        if ws:
            try:
                await ws.send(text)
                message_sent()
            except Exception:
                logger.exception("Send to %s failed", conn_id)