swiftpy-framework 0.1.0__py3-none-any.whl

swiftpy/websockets/connection.py

@@ -0,0 +1,317 @@
+from __future__ import annotations
+
+import asyncio
+import enum
+import json
+from collections.abc import Awaitable, Callable
+from typing import Any, cast
+from urllib.parse import parse_qs
+
+from swiftpy.http.response import encode_json
+
+
+class WebSocketState(enum.Enum):
+    CONNECTING = 0
+    CONNECTED = 1
+    DISCONNECTED = 2
+
+
+class WebSocketException(Exception):
+    """Base exception for WebSocket errors."""
+
+
+class WebSocketDisconnect(WebSocketException):
+    """Raised when a WebSocket connection is disconnected."""
+
+    def __init__(self, code: int = 1000, reason: str = "") -> None:
+        self.code = code
+        self.reason = reason
+        super().__init__(f"WebSocket disconnected with code {code} ({reason or 'no reason given'})")
+
+
+class WebSocketConnection:
+    """
+    Wraps an ASGI websocket connection.
+    Provides standard accept(), receive(), send(), and close() methods.
+    """
+
+    def __init__(
+        self,
+        scope: dict[str, Any],
+        receive: Callable[[], Awaitable[dict[str, Any]]],
+        send: Callable[[dict[str, Any]], Awaitable[None]],
+        max_message_size: int = 1_048_576,
+        idle_timeout: float | None = None,
+    ) -> None:
+        self.scope = scope
+        self._receive = receive
+        self._send = send
+        self.max_message_size = max_message_size
+        self.idle_timeout = idle_timeout
+
+        self.client_state = WebSocketState.CONNECTING
+        self.close_code: int | None = None
+        self.close_reason: str = ""
+        self.state: dict[str, Any] = {}
+        self._headers: dict[str, str] | None = None
+        self._query_param_list: dict[str, list[str]] | None = None
+        self._query_params: dict[str, str] | None = None
+
+    @property
+    def origin(self) -> str | None:
+        """The Origin header sent by the client, or None if absent."""
+        return self.headers.get("origin")
+
+    @property
+    def path(self) -> str:
+        """The URL path of the WebSocket request."""
+        return cast(str, self.scope.get("path", ""))
+
+    @property
+    def headers(self) -> dict[str, str]:
+        """Parsed request headers mapping lowercase header names to values (cached per connection)."""
+        cached = getattr(self, "_headers", None)
+        if cached is None:
+            result: dict[str, str] = {}
+            for k, v in self.scope.get("headers", []):
+                key = k.decode("latin-1").lower() if isinstance(k, bytes) else str(k).lower()
+                val = v.decode("latin-1") if isinstance(v, bytes) else str(v)
+                result[key] = val
+            cached = result
+            self._headers = cached
+        return cached
+
+    @property
+    def query_param_list(self) -> dict[str, list[str]]:
+        """Parsed dictionary of UTF-8 URL query parameters with all values as lists (cached)."""
+        cached = getattr(self, "_query_param_list", None)
+        if cached is None:
+            raw_qs = self.scope.get("query_string", b"")
+            if not raw_qs:
+                cached = {}
+            else:
+                qs_str = (
+                    raw_qs.decode("utf-8", errors="replace")
+                    if isinstance(raw_qs, bytes)
+                    else str(raw_qs)
+                )
+                cached = parse_qs(qs_str, keep_blank_values=True)
+            self._query_param_list = cached
+        return cached
+
+    @property
+    def query_params(self) -> dict[str, str]:
+        """Parsed dictionary of UTF-8 URL query string parameters (first value per key, cached)."""
+        cached = getattr(self, "_query_params", None)
+        if cached is None:
+            cached = {k: v[0] for k, v in self.query_param_list.items() if v}
+            self._query_params = cached
+        return cached
+
+    @property
+    def client(self) -> tuple[str, int] | None:
+        """The client's (host, port) tuple if available."""
+        client_tuple = self.scope.get("client")
+        if client_tuple and len(client_tuple) == 2:
+            return (str(client_tuple[0]), int(client_tuple[1]))
+        return None
+
+    @property
+    def subprotocols(self) -> list[str]:
+        """Subprotocols requested by the client."""
+        return list(self.scope.get("subprotocols", []))
+
+    @property
+    def is_connected(self) -> bool:
+        """Whether the WebSocket connection is actively open."""
+        return self.client_state == WebSocketState.CONNECTED
+
+    async def accept(
+        self,
+        subprotocol: str | None = None,
+        headers: list[tuple[bytes, bytes]] | None = None,
+    ) -> None:
+        """
+        Accept the incoming WebSocket connection handshake.
+        Consumes the initial 'websocket.connect' ASGI message and transmits 'websocket.accept'.
+        """
+        if self.client_state == WebSocketState.CONNECTED:
+            raise WebSocketException("Cannot accept an already accepted WebSocket connection")
+        if self.client_state == WebSocketState.DISCONNECTED:
+            raise WebSocketException("Cannot accept a closed WebSocket connection")
+
+        message = await self._receive()
+        msg_type = message.get("type")
+
+        if msg_type == "websocket.disconnect":
+            self.client_state = WebSocketState.DISCONNECTED
+            self.close_code = message.get("code", 1006)
+            self.close_reason = message.get("reason", "")
+            raise WebSocketDisconnect(code=self.close_code, reason=self.close_reason)
+
+        if msg_type != "websocket.connect":
+            raise WebSocketException(
+                f"Expected 'websocket.connect' message, but received '{msg_type}'"
+            )
+
+        accept_msg: dict[str, Any] = {"type": "websocket.accept"}
+        if subprotocol is not None:
+            accept_msg["subprotocol"] = subprotocol
+        if headers is not None:
+            accept_msg["headers"] = headers
+
+        await self._send(accept_msg)
+        self.client_state = WebSocketState.CONNECTED
+
+    async def receive(self) -> str | bytes:
+        """
+        Receive raw text or bytes from the WebSocket client.
+        Raises WebSocketDisconnect when the client disconnects, times out, or exceeds size limits.
+        """
+        if self.client_state == WebSocketState.CONNECTING:
+            raise WebSocketException("Cannot receive data: WebSocket has not been accepted yet")
+        if self.client_state == WebSocketState.DISCONNECTED:
+            raise WebSocketDisconnect(
+                code=self.close_code or 1000,
+                reason=self.close_reason,
+            )
+
+        try:
+            if self.idle_timeout is not None:
+                message = await asyncio.wait_for(self._receive(), timeout=self.idle_timeout)
+            else:
+                message = await self._receive()
+        except TimeoutError:
+            self.client_state = WebSocketState.DISCONNECTED
+            self.close_code = 1000
+            self.close_reason = "Idle timeout"
+            try:
+                await self._send({"type": "websocket.close", "code": 1000, "reason": "Idle timeout"})
+            except Exception:  # noqa: BLE001, S110
+                pass
+            raise WebSocketDisconnect(code=1000, reason="Idle timeout") from None
+
+        msg_type = message.get("type")
+
+        if msg_type == "websocket.receive":
+            raw_text = message.get("text")
+            raw_bytes = message.get("bytes")
+
+            if raw_text is not None:
+                text_str = cast(str, raw_text)
+                if len(text_str.encode("utf-8")) > self.max_message_size:
+                    self.client_state = WebSocketState.DISCONNECTED
+                    self.close_code = 1009
+                    self.close_reason = "Message too large"
+                    try:
+                        await self._send({"type": "websocket.close", "code": 1009, "reason": "Message too large"})
+                    except Exception:  # noqa: BLE001, S110
+                        pass
+                    raise WebSocketDisconnect(code=1009, reason="Message too large")
+                return text_str
+
+            if raw_bytes is not None:
+                bytes_val = cast(bytes, raw_bytes)
+                if len(bytes_val) > self.max_message_size:
+                    self.client_state = WebSocketState.DISCONNECTED
+                    self.close_code = 1009
+                    self.close_reason = "Message too large"
+                    try:
+                        await self._send({"type": "websocket.close", "code": 1009, "reason": "Message too large"})
+                    except Exception:  # noqa: BLE001, S110
+                        pass
+                    raise WebSocketDisconnect(code=1009, reason="Message too large")
+                return bytes_val
+
+            return ""
+
+        if msg_type == "websocket.disconnect":
+            self.client_state = WebSocketState.DISCONNECTED
+            self.close_code = message.get("code", 1000)
+            self.close_reason = message.get("reason", "")
+            raise WebSocketDisconnect(code=self.close_code, reason=self.close_reason)
+
+        raise WebSocketException(f"Unexpected WebSocket message type: '{msg_type}'")
+
+    async def receive_text(self) -> str:
+        """Receive a text message from the WebSocket client."""
+        data = await self.receive()
+        if not isinstance(data, str):
+            raise WebSocketException("Expected text frame, but received binary frame")
+        return data
+
+    async def receive_bytes(self) -> bytes:
+        """Receive a binary message from the WebSocket client."""
+        data = await self.receive()
+        if not isinstance(data, bytes):
+            raise WebSocketException("Expected binary frame, but received text frame")
+        return data
+
+    async def receive_json(self) -> Any:
+        """Receive and parse a JSON message from the WebSocket client."""
+        text = await self.receive_text()
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError as exc:
+            raise WebSocketException(f"Malformed JSON received: {exc.msg}") from exc
+
+    async def send(self, data: str | bytes | dict[str, Any] | list[Any]) -> None:
+        """
+        Send data to the WebSocket client.
+        Automatically routes str, bytes, and JSON serializable dict/list.
+        """
+        if self.client_state != WebSocketState.CONNECTED:
+            raise WebSocketException("Cannot send message: WebSocket is not connected")
+
+        if isinstance(data, str):
+            await self.send_text(data)
+        elif isinstance(data, bytes):
+            await self.send_bytes(data)
+        elif isinstance(data, (dict, list)):
+            await self.send_json(data)
+        else:
+            await self.send_text(str(data))
+
+    async def send_text(self, text: str) -> None:
+        """Send a text frame to the WebSocket client."""
+        if self.client_state != WebSocketState.CONNECTED:
+            raise WebSocketException("Cannot send message: WebSocket is not connected")
+        await self._send({"type": "websocket.send", "text": text})
+
+    async def send_bytes(self, data: bytes) -> None:
+        """Send a binary frame to the WebSocket client."""
+        if self.client_state != WebSocketState.CONNECTED:
+            raise WebSocketException("Cannot send message: WebSocket is not connected")
+        await self._send({"type": "websocket.send", "bytes": data})
+
+    async def send_json(self, data: Any) -> None:
+        """Serialize data to JSON using SwiftPY's strict encoder and send as a text frame."""
+        if self.client_state != WebSocketState.CONNECTED:
+            raise WebSocketException("Cannot send message: WebSocket is not connected")
+        await self._send(
+            {
+                "type": "websocket.send",
+                "text": encode_json(data).decode("utf-8"),
+            }
+        )
+
+    async def close(self, code: int = 1000, reason: str = "") -> None:
+        """Close the WebSocket connection."""
+        if self.client_state == WebSocketState.DISCONNECTED:
+            return
+
+        self.client_state = WebSocketState.DISCONNECTED
+        self.close_code = code
+        self.close_reason = reason
+        await self._send({"type": "websocket.close", "code": code, "reason": reason})
+
+
+WebSocket = WebSocketConnection
+
+__all__ = [
+    "WebSocket",
+    "WebSocketConnection",
+    "WebSocketDisconnect",
+    "WebSocketException",
+    "WebSocketState",
+]

swiftpy/websockets/manager.py

@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import uuid
+from collections import defaultdict
+from collections.abc import Callable
+from typing import Any
+from weakref import WeakSet
+
+from swiftpy.core.container import get_container
+from swiftpy.events.redis_bus import RedisEventBus
+from swiftpy.http.response import encode_json
+from swiftpy.websockets.connection import (
+    WebSocketConnection,
+    WebSocketDisconnect,
+    WebSocketException,
+)
+
+logger = logging.getLogger("swiftpy.websockets")
+
+
+class WebSocketManager:
+    """
+    Manages active WebSocket connections.
+    Supports rooms, filtering, and concurrent broadcast distribution.
+    """
+
+    _instance: WebSocketManager | None = None
+
+    def __init__(self) -> None:
+        self._connections: set[WebSocketConnection] = set()
+        self._rooms: dict[str, set[WebSocketConnection]] = defaultdict(set)
+        self._connection_rooms: dict[WebSocketConnection, set[str]] = defaultdict(set)
+        self._lock = asyncio.Lock()
+        self._instance_id: str = uuid.uuid4().hex
+        self._subscribed_buses: WeakSet[RedisEventBus] = WeakSet()
+
+    def subscribe_to_bus(self, bus: RedisEventBus) -> None:
+        """Subscribe this manager to distributed WebSocket broadcasts on the given RedisEventBus."""
+        if bus in self._subscribed_buses:
+            return
+        bus.subscribe("websocket.broadcast", self._handle_distributed_broadcast)
+        self._subscribed_buses.add(bus)
+
+    def unsubscribe_from_bus(self, bus: RedisEventBus) -> None:
+        """Unsubscribe this manager from distributed WebSocket broadcasts on the given RedisEventBus."""
+        if bus in self._subscribed_buses:
+            bus.unsubscribe("websocket.broadcast", self._handle_distributed_broadcast)
+            self._subscribed_buses.discard(bus)
+
+    async def _handle_distributed_broadcast(self, payload: Any) -> None:
+        """Deliver a cross-worker broadcast event locally while ignoring messages sent by this worker."""
+        if not isinstance(payload, dict):
+            return
+        if payload.get("sender_id") == self._instance_id:
+            return
+        if "message" not in payload:
+            return
+        room = payload.get("room")
+        await self.broadcast(
+            payload["message"],
+            room=room if isinstance(room, str) or room is None else str(room),
+            distribute=False,
+        )
+
+    @classmethod
+    def get_instance(cls) -> WebSocketManager:
+        """Retrieve the active application's WebSocketManager from the container, or fallback singleton."""
+        try:
+            container = get_container()
+            if container.has(cls):
+                return container.resolve(cls)
+        except RuntimeError:
+            pass
+
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def reset_instance(cls) -> None:
+        """Reset the fallback singleton instance and clear the active container's manager if present."""
+        cls._instance = None
+        try:
+            container = get_container()
+            if container.has(cls):
+                fresh = cls()
+                container.instance(cls, fresh)
+        except RuntimeError:
+            pass
+
+    @property
+    def active_count(self) -> int:
+        """The total number of currently active connections."""
+        return len(self._connections)
+
+    def is_connected(self, connection: WebSocketConnection) -> bool:
+        """Check if a specific connection is actively registered."""
+        return connection in self._connections
+
+    def get_connections(self, room: str | None = None) -> list[WebSocketConnection]:
+        """
+        Get a list of active connections.
+        If a room is specified, returns active connections in that room.
+        """
+        if room is not None:
+            return list(self._rooms.get(room, set()))
+        return list(self._connections)
+
+    async def connect(
+        self,
+        connection: WebSocketConnection,
+        room: str | None = None,
+    ) -> None:
+        """
+        Register an active WebSocket connection.
+        Optionally assign the connection to a room.
+        """
+        async with self._lock:
+            self._connections.add(connection)
+            if room is not None:
+                self._rooms[room].add(connection)
+                self._connection_rooms[connection].add(room)
+
+    async def disconnect(
+        self,
+        connection: WebSocketConnection,
+        room: str | None = None,
+    ) -> None:
+        """
+        Unregister a WebSocket connection.
+        Removes the connection from active connections and all rooms.
+        """
+        async with self._lock:
+            self._connections.discard(connection)
+            if connection in self._connection_rooms:
+                for r in list(self._connection_rooms[connection]):
+                    if r in self._rooms:
+                        self._rooms[r].discard(connection)
+                        if not self._rooms[r]:
+                            del self._rooms[r]
+                del self._connection_rooms[connection]
+            if room is not None and room in self._rooms:
+                self._rooms[room].discard(connection)
+                if not self._rooms[room]:
+                    del self._rooms[room]
+
+    async def join_room(self, room: str, connection: WebSocketConnection) -> None:
+        """Add an already connected WebSocket to a specific room."""
+        async with self._lock:
+            self._connections.add(connection)
+            self._rooms[room].add(connection)
+            self._connection_rooms[connection].add(room)
+
+    async def leave_room(self, room: str, connection: WebSocketConnection) -> None:
+        """Remove a WebSocket from a specific room without disconnecting it."""
+        async with self._lock:
+            if room in self._rooms:
+                self._rooms[room].discard(connection)
+                if not self._rooms[room]:
+                    del self._rooms[room]
+            if connection in self._connection_rooms:
+                self._connection_rooms[connection].discard(room)
+                if not self._connection_rooms[connection]:
+                    del self._connection_rooms[connection]
+
+    async def broadcast(
+        self,
+        message: str | bytes | dict[str, Any] | list[Any],
+        *,
+        room: str | None = None,
+        filter_fn: Callable[[WebSocketConnection], bool] | None = None,
+        timeout: float = 5.0,
+        max_concurrency: int = 100,
+        distribute: bool = False,
+    ) -> None:
+        """
+        Broadcast a message to active connections.
+        - If `room` is provided, delivers only to connections in that room.
+        - If `filter_fn` is provided, delivers only to connections matching the predicate.
+        - If neither is provided, delivers to all active connections.
+        - `timeout`: Maximum seconds to wait per connection send (prevents stalled clients from delaying others).
+        - `max_concurrency`: Bounded concurrency using asyncio.Semaphore.
+        - `distribute`: If True, also publishes the broadcast across workers via RedisEventBus.
+        Failed, timed out, or disconnected connections are automatically cleaned up.
+        """
+        async with self._lock:
+            if room is not None:
+                candidates = list(self._rooms.get(room, set()))
+            else:
+                candidates = list(self._connections)
+
+        if filter_fn is not None:
+            recipients = []
+            for conn in candidates:
+                try:
+                    if filter_fn(conn):
+                        recipients.append(conn)
+                except Exception as exc:  # noqa: BLE001
+                    logger.warning("Error evaluating filter_fn on WebSocket connection: %s", exc)
+        else:
+            recipients = candidates
+
+        # 1. Single serialization for JSON and text payloads (P9-W3, WS-2)
+        prepared_text: str | None = None
+        prepared_bytes: bytes | None = None
+        distributed_message: Any = message
+
+        if isinstance(message, (dict, list)):
+            prepared_text = encode_json(message).decode("utf-8")
+            distributed_message = json.loads(prepared_text)
+        elif isinstance(message, str):
+            prepared_text = message
+        elif isinstance(message, bytes):
+            prepared_bytes = message
+        else:
+            prepared_text = str(message)
+
+        if recipients:
+            sem = asyncio.Semaphore(max(1, max_concurrency))
+
+            async def _safe_send(conn: WebSocketConnection) -> tuple[WebSocketConnection, Exception | None]:
+                try:
+                    async with sem:
+                        if prepared_text is not None:
+                            await asyncio.wait_for(conn.send_text(prepared_text), timeout=timeout)
+                        elif prepared_bytes is not None:
+                            await asyncio.wait_for(conn.send_bytes(prepared_bytes), timeout=timeout)
+                    return conn, None
+                except (WebSocketDisconnect, WebSocketException) as exc:
+                    return conn, exc
+                except TimeoutError as exc:
+                    logger.warning(
+                        "Timeout (%ss) broadcasting to WebSocket connection; evicting stalled connection",
+                        timeout,
+                    )
+                    return conn, exc
+                except Exception as exc:  # noqa: BLE001
+                    # Failures on one client must not abort broadcast to remaining clients
+                    logger.warning("Error broadcasting to WebSocket connection: %s", exc)
+                    return conn, exc
+
+            results = await asyncio.gather(*[_safe_send(conn) for conn in recipients])
+
+            failed_conns = [conn for conn, exc in results if exc is not None]
+            if failed_conns:
+                for conn in failed_conns:
+                    await self.disconnect(conn)
+
+        # 2. Cross-worker distributed fan-out via RedisEventBus (P9-W3, P6-Q6, WS-2)
+        if distribute:
+            try:
+                container = get_container()
+                if container.has(RedisEventBus):
+                    bus = container.resolve(RedisEventBus)
+                    await bus.publish(
+                        "websocket.broadcast",
+                        {
+                            "message": distributed_message,
+                            "room": room,
+                            "sender_id": self._instance_id,
+                        },
+                        local=False,
+                    )
+            except Exception as exc:  # noqa: BLE001
+                logger.warning("Error publishing distributed WebSocket broadcast: %s", exc)
+
+
+__all__ = ["WebSocketManager"]