orchestrata_ultimate.cxp.ws_server

Reference › orchestrata_ultimate › cxp › ws_server

WSServer is a WebSocket server for CXP, an alternative to the cxpt direct transport that the rest of the system uses. No other module in the project imports it, and it is separate from the gateway's own WebSocket API. WSServerConfig holds host, port, maximum message size, and ping interval and timeout. CXPWebSocketHandler serves each connection: _process_message parses JSON text and dispatches on a type field. send_packet calls _handle_send_packet, subscribe calls _handle_subscribe, acknowledge calls _handle_ack, and ping gets a pong. An unknown type, invalid JSON or an exception produces an error message. WSServer wraps the handler: start begins serving with the configured limits, broadcast sends a packet to connected clients, stop shuts down, and serve_forever blocks. The matching client is ws_client. See cxp.

At a glance

Class diagram

Dependency diagram

No internal imports and no importers.

API

Classes

WSServerConfig (dataclass)

Decorators: dataclass · Line 30

Configuration for CXP WebSocket server.

CXPWebSocketHandler (dataclass)

Decorators: dataclass · Line 41

Handler for CXP WebSocket messages.

WSServer

Line 182

CXP WebSocket Server.

Source

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

# mypy: ignore-errors
"""WebSocket Server for CXP Protocol.

Provides a WebSocket-based transport layer for CXP packets,
enabling real-time inter-Volk communication.

Usage:
    server = WSServer(host="0.0.0.0", port=8765)
    await server.start()
    await server.serve_forever()
"""

import asyncio
import json
import logging
import uuid
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

import websockets

if TYPE_CHECKING:
    from websockets.asyncio.server import Server, ServerConnection

logger = logging.getLogger(__name__)


@dataclass
class WSServerConfig:
    """Configuration for CXP WebSocket server."""

    host: str = "0.0.0.0"
    port: int = 8765
    max_size: int = 10 * 1024 * 1024  # 10MB
    ping_interval: int = 30
    ping_timeout: int = 10


@dataclass
class CXPWebSocketHandler:
    """Handler for CXP WebSocket messages."""

    packet_handler: Callable[[dict[str, Any]], Any] | None = None
    subscriptions: dict[str, set[str]] = field(default_factory=dict)

    async def handle(self, websocket: ServerConnection, path: str = "/") -> None:
        """Handle WebSocket connection."""
        client_id = str(uuid.uuid4())
        logger.info(
            f"Client {client_id} connected from {websocket.remote_address} to path {path}"
        )

        try:
            await websocket.send(
                json.dumps(
                    {
                        "type": "connected",
                        "client_id": client_id,
                    }
                )
            )

            async for message in websocket:
                await self._process_message(websocket, message, client_id)

        except websockets.exceptions.ConnectionClosed:
            logger.info(f"Client {client_id} disconnected")
        except Exception as e:
            logger.error(f"Error handling client {client_id}: {e}")
        finally:
            if client_id in self.subscriptions:
                del self.subscriptions[client_id]

    async def _process_message(
        self,
        websocket: ServerConnection,
        message: str | bytes,
        client_id: str,
    ) -> None:
        """Process incoming WebSocket message."""
        try:
            data = json.loads(message)
            msg_type = data.get("type")

            if msg_type == "send_packet":
                await self._handle_send_packet(websocket, data)
            elif msg_type == "subscribe":
                await self._handle_subscribe(websocket, data, client_id)
            elif msg_type == "acknowledge":
                await self._handle_ack(websocket, data)
            elif msg_type == "ping":
                await websocket.send(json.dumps({"type": "pong"}))
            else:
                await websocket.send(
                    json.dumps(
                        {
                            "type": "error",
                            "message": f"Unknown message type: {msg_type}",
                        }
                    )
                )

        except json.JSONDecodeError:
            await websocket.send(
                json.dumps(
                    {
                        "type": "error",
                        "message": "Invalid JSON",
                    }
                )
            )
        except Exception as e:
            await websocket.send(
                json.dumps(
                    {
                        "type": "error",
                        "message": str(e),
                    }
                )
            )

    async def _handle_send_packet(
        self, websocket: ServerConnection, data: dict
    ) -> None:
        """Handle send_packet message."""
        packet = data.get("packet", {})

        if self.packet_handler:
            handler_result = self.packet_handler(packet)
            if asyncio.iscoroutine(handler_result) or hasattr(
                handler_result, "__await__"
            ):
                result = await handler_result
            else:
                result = handler_result
        else:
            result = {"success": True}

        await websocket.send(
            json.dumps(
                {
                    "type": "packet_sent",
                    "message_id": packet.get("message_id"),
                    "success": result.get("success", True),
                    "error": result.get("error", ""),
                }
            )
        )

    async def _handle_subscribe(
        self,
        websocket: ServerConnection,
        data: dict,
        client_id: str,
    ) -> None:
        """Handle subscribe message."""
        subjects = data.get("subjects", [])
        self.subscriptions[client_id] = set(subjects)

        await websocket.send(
            json.dumps(
                {
                    "type": "subscribed",
                    "subjects": subjects,
                }
            )
        )

    async def _handle_ack(self, websocket: ServerConnection, data: dict) -> None:
        """Handle acknowledge message."""
        await websocket.send(
            json.dumps(
                {
                    "type": "acknowledged",
                    "message_id": data.get("message_id"),
                }
            )
        )


class WSServer:
    """CXP WebSocket Server."""

    def __init__(
        self,
        config: WSServerConfig | None = None,
        packet_handler: Callable[[dict[str, Any]], Any] | None = None,
    ) -> None:
        self._config = config or WSServerConfig()
        self._handler = CXPWebSocketHandler(packet_handler=packet_handler)
        self._server: Server | None = None

    async def start(self) -> None:
        """Start the WebSocket server."""

        async def handler_wrapper(websocket: ServerConnection, path: str = "/") -> None:
            await self._handler.handle(websocket, path)

        self._server = await websockets.serve(
            handler_wrapper,
            self._config.host,
            self._config.port,
            max_size=self._config.max_size,
            ping_interval=self._config.ping_interval,
            ping_timeout=self._config.ping_timeout,
        )
        logger.info(
            f"CXP WebSocket server started on {self._config.host}:{self._config.port}"
        )

    async def stop(self) -> None:
        """Stop the WebSocket server."""
        if self._server:
            self._server.close()
            await self._server.wait_closed()
            logger.info("CXP WebSocket server stopped")

    async def broadcast(self, packet: dict[str, Any]) -> None:
        """Broadcast packet to all connected clients."""
        if not self._server:
            return

        message = json.dumps(
            {
                "type": "packet",
                "packet": packet,
            }
        )

        for ws in self._server.connections:
            await ws.send(message)

    async def serve_forever(self) -> None:
        """Run the server forever."""
        await asyncio.Future()