kryten-robot 0.6.9__py3-none-any.whl

kryten/nats_client.py

@@ -0,0 +1,468 @@
+"""NATS Client Wrapper for Event Publishing and Subscriptions.
+
+This module provides an asynchronous NATS client wrapper with connection
+lifecycle management, automatic reconnection, and error handling.
+"""
+
+import asyncio
+import logging
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from nats.aio.client import Client as NATS
+from nats.aio.subscription import Subscription
+
+from .config import NatsConfig
+from .errors import ConnectionError, NotConnectedError
+
+
+class NatsClient:
+    """Asynchronous NATS client wrapper with lifecycle management.
+
+    Provides a clean interface for connecting to NATS servers and publishing
+    messages with automatic reconnection and error handling.
+
+    Attributes:
+        config: NATS connection configuration.
+        logger: Logger instance for structured logging.
+        is_connected: Whether currently connected to NATS.
+        stats: Publishing statistics (messages, bytes, errors).
+
+    Examples:
+        >>> config = NatsConfig(servers=["nats://localhost:4222"])
+        >>> async with NatsClient(config, logger) as client:
+        ...     await client.publish("test.subject", b"hello")
+    """
+
+    def __init__(self, config: NatsConfig, logger: logging.Logger):
+        """Initialize NATS client.
+
+        Args:
+            config: NATS connection configuration.
+            logger: Logger for structured output.
+        """
+        self.config = config
+        self.logger = logger
+        self._nc: NATS | None = None
+        self._connected = False
+
+        # Connection tracking
+        self._connected_since: float | None = None
+        self._reconnect_count: int = 0
+
+        # Statistics tracking
+        self._messages_published = 0
+        self._bytes_sent = 0
+        self._errors = 0
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if currently connected to NATS.
+
+        Returns:
+            True if connected and client is active, False otherwise.
+        """
+        return self._connected and self._nc is not None and self._nc.is_connected
+
+    @property
+    def connected_since(self) -> float | None:
+        """Get timestamp when connection was established.
+
+        Returns:
+            Unix timestamp of connection time, or None if not connected.
+        """
+        return self._connected_since
+
+    @property
+    def reconnect_count(self) -> int:
+        """Get number of reconnection attempts.
+
+        Returns:
+            Count of reconnections since instance creation.
+        """
+        return self._reconnect_count
+
+    @property
+    def connected_url(self) -> str | None:
+        """Get the currently connected NATS server URL.
+
+        Returns:
+            Server URL if connected, None otherwise.
+        """
+        if self._nc and self._nc.is_connected:
+            return self._nc.connected_url.netloc
+        return None
+
+    @property
+    def stats(self) -> dict[str, int]:
+        """Get publishing statistics.
+
+        Returns:
+            Dictionary with 'messages_published', 'bytes_sent', 'errors' counts.
+
+        Examples:
+            >>> stats = client.stats
+            >>> print(f"Published: {stats['messages_published']}")
+        """
+        return {
+            'messages_published': self._messages_published,
+            'bytes_sent': self._bytes_sent,
+            'errors': self._errors,
+        }
+
+    async def connect(self) -> None:
+        """Establish connection to NATS servers.
+
+        Attempts to connect to configured NATS servers with automatic
+        reconnection support if enabled.
+
+        Raises:
+            ConnectionError: If connection cannot be established.
+            asyncio.CancelledError: If connection is cancelled.
+
+        Examples:
+            >>> client = NatsClient(config, logger)
+            >>> await client.connect()
+            >>> assert client.is_connected
+        """
+        if self._connected:
+            self.logger.warning("Already connected to NATS, ignoring connect() call")
+            return
+
+        self.logger.info(
+            "Connecting to NATS",
+            extra={"servers": self.config.servers}
+        )
+
+        try:
+            # Create NATS client
+            self._nc = NATS()
+
+            # Configure connection options
+            options = {
+                "servers": self.config.servers,
+                "connect_timeout": self.config.connect_timeout,
+                "max_reconnect_attempts": self.config.max_reconnect_attempts if self.config.allow_reconnect else 0,
+                "reconnect_time_wait": self.config.reconnect_time_wait,
+                "allow_reconnect": self.config.allow_reconnect,
+                "error_cb": self._error_callback,
+                "disconnected_cb": self._disconnected_callback,
+                "reconnected_cb": self._reconnected_callback,
+                "closed_cb": self._closed_callback,
+            }
+
+            # Add credentials if provided
+            if self.config.user and self.config.password:
+                options["user"] = self.config.user
+                options["password"] = self.config.password
+                self.logger.debug("Using NATS credentials for authentication")
+
+            # Connect
+            await self._nc.connect(**options)
+
+            self._connected = True
+
+            # Track connection timing
+            import time
+            self._connected_since = time.time()
+
+            self.logger.info(
+                "Connected to NATS",
+                extra={
+                    "servers": self.config.servers,
+                    "server_info": self._nc.connected_server_version
+                }
+            )
+
+        except asyncio.CancelledError:
+            self.logger.warning("NATS connection cancelled")
+            await self._cleanup()
+            raise
+
+        except Exception as e:
+            self.logger.error(
+                "Failed to connect to NATS",
+                extra={"servers": self.config.servers, "error": str(e)}
+            )
+            await self._cleanup()
+            raise ConnectionError(f"Failed to connect to NATS: {e}") from e
+
+    async def disconnect(self) -> None:
+        """Drain and close NATS connection gracefully.
+
+        Flushes pending messages before disconnecting. Safe to call
+        multiple times or when not connected.
+
+        Examples:
+            >>> await client.connect()
+            >>> await client.disconnect()
+            >>> assert not client.is_connected
+        """
+        if not self._connected:
+            self.logger.debug("Not connected to NATS, disconnect() is a no-op")
+            return
+
+        self.logger.info("Disconnecting from NATS")
+
+        await self._cleanup()
+
+        self.logger.info("Disconnected from NATS")
+
+    async def _cleanup(self) -> None:
+        """Clean up resources and reset state.
+
+        Internal method for draining, closing connection, and resetting state.
+        """
+        self._connected = False
+        self._connected_since = None
+
+        if self._nc:
+            try:
+                # Drain to flush pending messages
+                await self._nc.drain()
+            except Exception as e:
+                self.logger.warning(f"Error draining NATS connection: {e}")
+
+            try:
+                # Close connection
+                if not self._nc.is_closed:
+                    await self._nc.close()
+            except Exception as e:
+                self.logger.warning(f"Error closing NATS connection: {e}")
+
+            finally:
+                self._nc = None
+
+    async def publish(self, subject: str, data: bytes) -> None:
+        """Publish message to NATS subject.
+
+        Args:
+            subject: NATS subject name (e.g., "cytube.events.chatMsg").
+            data: Message payload as bytes.
+
+        Raises:
+            NotConnectedError: If not connected to NATS.
+            ValueError: If subject is empty or data is not bytes.
+
+        Examples:
+            >>> await client.publish("test.subject", b"hello world")
+        """
+        if not self.is_connected:
+            raise NotConnectedError("Not connected to NATS server")
+
+        if not subject:
+            raise ValueError("Subject cannot be empty")
+
+        if not isinstance(data, bytes):
+            raise ValueError("Data must be bytes")
+
+        try:
+            await self._nc.publish(subject, data)
+            self._messages_published += 1
+            self._bytes_sent += len(data)
+
+            self.logger.debug(
+                "Published message to NATS",
+                extra={"subject": subject, "size": len(data)}
+            )
+
+        except Exception as e:
+            self._errors += 1
+            self.logger.error(
+                "Failed to publish message",
+                extra={"subject": subject, "error": str(e)}
+            )
+            raise
+
+    async def subscribe(
+        self,
+        subject: str,
+        callback: Callable[[str, bytes], Awaitable[None]]
+    ) -> Subscription:
+        """Subscribe to NATS subject with callback.
+
+        Args:
+            subject: NATS subject pattern (e.g., "cytube.commands.>").
+            callback: Async callback function(subject: str, data: bytes).
+
+        Returns:
+            Subscription object that can be used to unsubscribe.
+
+        Raises:
+            NotConnectedError: If not connected to NATS.
+            ValueError: If subject is empty or callback is not callable.
+
+        Examples:
+            >>> async def handler(subject: str, data: bytes):
+            ...     print(f"Received on {subject}: {data}")
+            >>> sub = await client.subscribe("test.>", handler)
+        """
+        if not self.is_connected:
+            raise NotConnectedError("Not connected to NATS server")
+
+        if not subject:
+            raise ValueError("Subject cannot be empty")
+
+        if not callable(callback):
+            raise ValueError("Callback must be callable")
+
+        try:
+            # NATS callback receives a Msg object, we need to adapt it
+            async def nats_callback(msg):
+                await callback(msg.subject, msg.data)
+
+            subscription = await self._nc.subscribe(subject, cb=nats_callback)
+
+            self.logger.debug(
+                "Subscribed to NATS subject",
+                extra={"subject": subject}
+            )
+
+            return subscription
+
+        except Exception as e:
+            self._errors += 1
+            self.logger.error(
+                "Failed to subscribe to subject",
+                extra={"subject": subject, "error": str(e)}
+            )
+            raise
+
+    async def subscribe_request_reply(
+        self,
+        subject: str,
+        callback: Callable[[Any], Awaitable[None]]
+    ) -> Subscription:
+        """Subscribe to NATS subject for request-reply pattern.
+
+        Unlike regular subscribe, this passes the full NATS Msg object to the callback,
+        allowing access to msg.reply for sending responses.
+
+        Args:
+            subject: NATS subject pattern (e.g., "kryten.robot.command").
+            callback: Async callback function(msg: Msg) that receives full message.
+
+        Returns:
+            Subscription object that can be used to unsubscribe.
+
+        Raises:
+            NotConnectedError: If not connected to NATS.
+            ValueError: If subject is empty or callback is not callable.
+
+        Examples:
+            >>> async def handler(msg):
+            ...     request = json.loads(msg.data.decode())
+            ...     response = {"result": "ok"}
+            ...     await nats.publish(msg.reply, json.dumps(response).encode())
+            >>> sub = await client.subscribe_request_reply("kryten.robot.command", handler)
+        """
+        if not self.is_connected:
+            raise NotConnectedError("Not connected to NATS server")
+
+        if not subject:
+            raise ValueError("Subject cannot be empty")
+
+        if not callable(callback):
+            raise ValueError("Callback must be callable")
+
+        try:
+            # Pass the full msg object directly to callback
+            subscription = await self._nc.subscribe(subject, cb=callback)
+
+            self.logger.debug(
+                "Subscribed to NATS subject (request-reply)",
+                extra={"subject": subject}
+            )
+
+            return subscription
+
+        except Exception as e:
+            self._errors += 1
+            self.logger.error(
+                "Failed to subscribe to subject",
+                extra={"subject": subject, "error": str(e)}
+            )
+            raise
+
+    async def unsubscribe(self, subscription: Subscription) -> None:
+        """Unsubscribe from NATS subject.
+
+        Args:
+            subscription: Subscription object from subscribe().
+
+        Examples:
+            >>> sub = await client.subscribe("test.>", handler)
+            >>> await client.unsubscribe(sub)
+        """
+        if subscription:
+            try:
+                await subscription.unsubscribe()
+                self.logger.debug("Unsubscribed from NATS subject")
+            except Exception as e:
+                self.logger.warning(f"Error unsubscribing: {e}")
+
+    async def _error_callback(self, error: Exception) -> None:
+        """Callback for NATS connection errors.
+
+        Args:
+            error: Exception that occurred.
+        """
+        self.logger.error(f"NATS error: {error}")
+        self._errors += 1
+
+    async def _disconnected_callback(self) -> None:
+        """Callback when disconnected from NATS."""
+        self.logger.warning("Disconnected from NATS server")
+        self._connected = False
+
+    async def _reconnected_callback(self) -> None:
+        """Callback when reconnected to NATS."""
+        self._reconnect_count += 1
+
+        # Update connected_since for the new connection
+        import time
+        self._connected_since = time.time()
+
+        self.logger.info(
+            f"Reconnected to NATS server (reconnect #{self._reconnect_count})"
+        )
+        self._connected = True
+
+    async def _closed_callback(self) -> None:
+        """Callback when NATS connection is closed."""
+        self.logger.info("NATS connection closed")
+        self._connected = False
+
+    async def __aenter__(self):
+        """Async context manager entry.
+
+        Automatically connects when entering the context.
+
+        Returns:
+            Self for use in context.
+
+        Examples:
+            >>> async with NatsClient(config, logger) as client:
+            ...     await client.publish("test", b"data")
+        """
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit.
+
+        Ensures disconnect is called even if an exception occurs.
+
+        Args:
+            exc_type: Exception type if raised.
+            exc_val: Exception value if raised.
+            exc_tb: Exception traceback if raised.
+
+        Returns:
+            False to propagate exceptions.
+        """
+        await self.disconnect()
+        return False
+
+
+__all__ = ["NatsClient"]

kryten/raw_event.py

@@ -0,0 +1,165 @@
+"""Raw Event Dataclass for CyTube Socket.IO Events.
+
+This module defines the RawEvent dataclass for wrapping unparsed CyTube events
+with metadata before publishing to NATS.
+"""
+
+import json
+import uuid
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+from typing import Any
+
+
+@dataclass(frozen=True)
+class RawEvent:
+    """Immutable container for raw Socket.IO events with metadata.
+
+    Wraps CyTube Socket.IO events with timestamps, correlation IDs, and channel
+    information for NATS publishing and distributed tracing.
+
+    Attributes:
+        event_name: Socket.IO event name (e.g., "chatMsg", "addUser").
+        payload: Raw Socket.IO event data as dictionary.
+        channel: CyTube channel name.
+        domain: CyTube server domain (e.g., "cytu.be").
+        timestamp: UTC ISO 8601 timestamp with microseconds.
+        correlation_id: UUID4 for distributed tracing.
+
+    Examples:
+        >>> event = RawEvent(
+        ...     event_name="chatMsg",
+        ...     payload={"user": "bob", "msg": "hello"},
+        ...     channel="lounge",
+        ...     domain="cytu.be"
+        ... )
+        >>> json_bytes = event.to_bytes()
+        >>> await nats_client.publish("cytube.events.chatMsg", json_bytes)
+    """
+
+    event_name: str
+    payload: dict[str, Any]
+    channel: str
+    domain: str
+    timestamp: str = field(
+        default_factory=lambda: datetime.now(UTC).isoformat()
+    )
+    correlation_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary representation.
+
+        Returns:
+            Dictionary with all fields as key-value pairs.
+
+        Examples:
+            >>> event = RawEvent("test", {}, "ch", "dom")
+            >>> d = event.to_dict()
+            >>> assert "event_name" in d
+            >>> assert "timestamp" in d
+        """
+        return {
+            "event_name": self.event_name,
+            "payload": self.payload,
+            "channel": self.channel,
+            "domain": self.domain,
+            "timestamp": self.timestamp,
+            "correlation_id": self.correlation_id,
+        }
+
+    def to_json(self) -> str:
+        """Serialize to JSON string.
+
+        Returns:
+            JSON string representation of the event.
+
+        Examples:
+            >>> event = RawEvent("chatMsg", {"msg": "hi"}, "ch", "dom")
+            >>> json_str = event.to_json()
+            >>> assert "chatMsg" in json_str
+        """
+        return json.dumps(self.to_dict())
+
+    def to_bytes(self) -> bytes:
+        """Serialize to UTF-8 encoded JSON bytes.
+
+        Returns:
+            UTF-8 encoded JSON bytes suitable for NATS publishing.
+
+        Examples:
+            >>> event = RawEvent("test", {}, "ch", "dom")
+            >>> data = event.to_bytes()
+            >>> assert isinstance(data, bytes)
+        """
+        return self.to_json().encode("utf-8")
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "RawEvent":
+        """Deserialize from dictionary.
+
+        Args:
+            data: Dictionary with event fields.
+
+        Returns:
+            RawEvent instance with fields from dictionary.
+
+        Examples:
+            >>> data = {
+            ...     "event_name": "test",
+            ...     "payload": {},
+            ...     "channel": "ch",
+            ...     "domain": "dom",
+            ...     "timestamp": "2024-01-15T10:30:00.123456Z",
+            ...     "correlation_id": "550e8400-e29b-41d4-a716-446655440000"
+            ... }
+            >>> event = RawEvent.from_dict(data)
+            >>> assert event.event_name == "test"
+        """
+        return cls(
+            event_name=data["event_name"],
+            payload=data["payload"],
+            channel=data["channel"],
+            domain=data["domain"],
+            timestamp=data.get(
+                "timestamp", datetime.now(UTC).isoformat()
+            ),
+            correlation_id=data.get("correlation_id", str(uuid.uuid4())),
+        )
+
+    def __str__(self) -> str:
+        """String representation for logging.
+
+        Returns:
+            Human-readable string with key event information.
+
+        Examples:
+            >>> event = RawEvent("chatMsg", {}, "lounge", "cytu.be")
+            >>> s = str(event)
+            >>> assert "chatMsg" in s
+            >>> assert "lounge" in s
+        """
+        return (
+            f"RawEvent(event={self.event_name}, channel={self.channel}, "
+            f"domain={self.domain}, id={self.correlation_id[:8]}...)"
+        )
+
+    def __repr__(self) -> str:
+        """Detailed representation for debugging.
+
+        Returns:
+            Full representation with all fields.
+
+        Examples:
+            >>> event = RawEvent("test", {}, "ch", "dom")
+            >>> repr(event)  # doctest: +ELLIPSIS
+            "RawEvent(event_name='test', ...)"
+        """
+        return (
+            f"RawEvent(event_name={self.event_name!r}, "
+            f"payload={self.payload!r}, channel={self.channel!r}, "
+            f"domain={self.domain!r}, timestamp={self.timestamp!r}, "
+            f"correlation_id={self.correlation_id!r})"
+        )
+
+
+__all__ = ["RawEvent"]