orchestrata_ultimate.cxp.ws_client

Reference › orchestrata_ultimate › cxp › ws_client

WSClient is a WebSocket client for CXP, the counterpart of ws_server. No other module in the project imports it. WSClientConfig holds the URI, which comes from CXP_WS_URI when the config object is constructed (not when the module is imported, so later environment changes are seen), plus a timeout and a maximum message size. connect opens the socket and starts _receive_loop, which passes each message to _process_message and queues the reply. send_packet connects if needed, sends a send_packet message, waits for the queued reply within the timeout and returns success, message id and error. subscribe registers subjects, receive_packets is an async generator of incoming packets, acknowledge confirms a message, ping checks the link, and disconnect closes it. See cxp.

At a glance

Class diagram

Dependency diagram

No internal imports and no importers.

API

Classes

WSClientConfig (dataclass)

Decorators: dataclass · Line 28

Configuration for CXP WebSocket client.

WSClient

Line 43

CXP WebSocket Client.

Source

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

"""WebSocket Client for CXP Protocol.

Provides a WebSocket-based client for sending and receiving CXP packets.

Usage:
    client = WSClient(uri="ws://localhost:8765")
    await client.connect()
    receipt = await client.send_packet(packet)
    async for packet in client.receive_packets():
        print(packet)
    await client.disconnect()
"""

import asyncio
import json
import logging
import os
from collections.abc import AsyncGenerator
from dataclasses import dataclass, field
from typing import Any

import websockets

logger = logging.getLogger(__name__)


@dataclass
class WSClientConfig:
    """Configuration for CXP WebSocket client.

    ``uri`` resolves from ``CXP_WS_URI`` at CONSTRUCTION time (a
    ``default_factory``, not an import-time read): a plain field default
    would bake whatever the environment held when this module was first
    imported, making later env changes (CLI entrypoints layering a profile,
    tests injecting an endpoint) silently invisible.
    """

    uri: str = field(default_factory=lambda: os.environ.get("CXP_WS_URI", ""))
    timeout_seconds: int = 30
    max_size: int = 10 * 1024 * 1024  # 10MB


class WSClient:
    """CXP WebSocket Client."""

    def __init__(
        self,
        config: WSClientConfig | None = None,
    ) -> None:
        self._config = config or WSClientConfig()
        self._websocket: websockets.ClientConnection | None = None
        self._connected = False
        self._client_id: str | None = None
        self._receive_task: asyncio.Task | None = None
        self._message_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self._subscriptions: set[str] = set()

    def _require_ws(self) -> websockets.ClientConnection:
        """Return the live connection, raising if not connected.

        ``self._websocket`` is ``None`` until ``connect()`` succeeds; every
        send/recv site below only runs after a ``connect()`` call (directly
        or via the ``if not self._connected: await self.connect()`` guard
        each public method already has), so this documents that invariant
        for pyright once, instead of an ``assert``/``# type: ignore`` at
        each of the 7 access sites.
        """
        if self._websocket is None:
            raise RuntimeError("WSClient is not connected — call connect() first")
        return self._websocket

    async def connect(self) -> None:
        """Connect to the WebSocket server."""
        self._websocket = await websockets.connect(
            self._config.uri,
            max_size=self._config.max_size,
        )
        self._connected = True
        logger.info(f"Connected to CXP WebSocket server at {self._config.uri}")

        response = await self._require_ws().recv()
        data = json.loads(response)
        if data.get("type") == "connected":
            self._client_id = data.get("client_id")
            logger.info(f"Assigned client ID: {self._client_id}")

        self._receive_task = asyncio.create_task(self._receive_loop())

    async def disconnect(self) -> None:
        """Disconnect from the WebSocket server."""
        if self._receive_task:
            self._receive_task.cancel()
            try:
                await self._receive_task
            except asyncio.CancelledError:
                logger.debug("receive task cancelled during disconnect")

        if self._websocket:
            await self._require_ws().close()
            self._connected = False
            logger.info("Disconnected from CXP WebSocket server")

    async def _receive_loop(self) -> None:
        """Receive messages from server."""
        try:
            async for message in self._require_ws():
                await self._process_message(message)
        except websockets.exceptions.ConnectionClosed:
            logger.info("Connection closed")
        except Exception as e:
            logger.error(f"Error in receive loop: {e}")

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

        if msg_type == "packet":
            await self._message_queue.put(data.get("packet", {}))
        elif msg_type in ("packet_sent", "acknowledged", "subscribed", "pong"):
            await self._message_queue.put(data)
        elif msg_type == "error":
            logger.error(f"Server error: {data.get('message')}")
            await self._message_queue.put(data)

    async def send_packet(self, packet: dict[str, Any]) -> dict[str, Any]:
        """Send a CXP packet via WebSocket."""
        if not self._connected:
            await self.connect()

        await self._require_ws().send(
            json.dumps(
                {
                    "type": "send_packet",
                    "packet": packet,
                }
            )
        )

        response = await asyncio.wait_for(
            self._message_queue.get(),
            timeout=self._config.timeout_seconds,
        )

        return {
            "success": response.get("success", False),
            "message_id": response.get("message_id"),
            "error": response.get("error", ""),
        }

    async def subscribe(self, subjects: list[str]) -> bool:
        """Subscribe to packets with specific subjects."""
        if not self._connected:
            await self.connect()

        await self._require_ws().send(
            json.dumps(
                {
                    "type": "subscribe",
                    "subjects": subjects,
                }
            )
        )

        self._subscriptions.update(subjects)

        response = await asyncio.wait_for(
            self._message_queue.get(),
            timeout=self._config.timeout_seconds,
        )

        return response.get("type") == "subscribed"

    async def receive_packets(self) -> AsyncGenerator[dict[str, Any]]:
        """Receive packets from server."""
        while self._connected:
            try:
                packet = await asyncio.wait_for(
                    self._message_queue.get(),
                    timeout=1.0,
                )
                if packet.get("type") != "packet":
                    continue
                yield packet
            except TimeoutError:
                continue
            except asyncio.CancelledError:
                break

    async def acknowledge(self, message_id: str) -> bool:
        """Acknowledge receipt of a packet."""
        if not self._connected:
            await self.connect()

        await self._require_ws().send(
            json.dumps(
                {
                    "type": "acknowledge",
                    "message_id": message_id,
                }
            )
        )

        response = await asyncio.wait_for(
            self._message_queue.get(),
            timeout=self._config.timeout_seconds,
        )

        return response.get("type") == "acknowledged"

    async def ping(self) -> bool:
        """Send ping to server."""
        if not self._connected:
            return False

        await self._require_ws().send(json.dumps({"type": "ping"}))

        try:
            response = await asyncio.wait_for(
                self._message_queue.get(),
                timeout=5.0,
            )
            return response.get("type") == "pong"
        except TimeoutError:
            return False