ruvnet--RuView/python/wifi_densepose/client/ws.py

"""ADR-117 P4 — Asyncio WebSocket client for the sensing-server.

The Rust sensing-server (`v2/crates/wifi-densepose-sensing-server`)
broadcasts three structured message types over `ws://<host>:<port>/ws/sensing`:

| `type` field | Source line in main.rs | Payload shape |
|---|---|---|
| `connection_established` | 2596 | `{node_id, version, capabilities}` |
| `pose_data` | 2655 | `{node_id, timestamp, persons: [...], confidence}` |
| `edge_vitals` | 4548 | `{node_id, presence, fall_detected, motion, breathing_rate_bpm, heartrate_bpm, ...}` |

`SensingClient` is a pure-Python asyncio wrapper around `websockets>=12`
that connects, decodes JSON, and yields typed dataclasses.

Example:

```python
import asyncio
from wifi_densepose.client import SensingClient, EdgeVitalsMessage

async def main():
    async with SensingClient("ws://localhost:8765/ws/sensing") as client:
        async for msg in client.stream():
            if isinstance(msg, EdgeVitalsMessage):
                print(f"BR={msg.breathing_rate_bpm}, HR={msg.heartrate_bpm}")

asyncio.run(main())
```
"""

from __future__ import annotations

import asyncio
import json
import logging
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Optional

# Defer import — only fail at construction time, not at module load.
try:
    import websockets  # type: ignore[import-not-found]
    from websockets.exceptions import ConnectionClosed  # type: ignore[import-not-found]
    _WEBSOCKETS_AVAILABLE = True
except ImportError:  # pragma: no cover
    _WEBSOCKETS_AVAILABLE = False


log = logging.getLogger(__name__)


# ─── Typed messages ──────────────────────────────────────────────────


@dataclass(frozen=True)
class SensingMessage:
    """Base class for typed sensing-server messages. The original JSON
    payload is preserved in ``raw`` for forward-compatibility with
    fields not yet modelled here."""
    type: str
    raw: dict[str, Any] = field(default_factory=dict, hash=False, compare=False)


@dataclass(frozen=True)
class ConnectionEstablishedMessage(SensingMessage):
    """First message after a successful WS handshake. Lets the client
    discover the node ID and capability flags without making a separate
    REST call."""
    node_id: str = ""
    version: str = ""
    capabilities: tuple[str, ...] = ()


@dataclass(frozen=True)
class EdgeVitalsMessage(SensingMessage):
    """Vital-sign telemetry fused from the edge-vitals path
    (ADR-021/ADR-110). Optional fields may be ``None`` when the
    upstream channel hasn't produced a measurement yet."""
    node_id: str = ""
    presence: bool = False
    fall_detected: bool = False
    motion: float = 0.0
    breathing_rate_bpm: Optional[float] = None
    heartrate_bpm: Optional[float] = None
    n_persons: int = 0
    motion_energy: float = 0.0
    presence_score: float = 0.0
    rssi: Optional[float] = None


@dataclass(frozen=True)
class PoseDataMessage(SensingMessage):
    """17-keypoint pose data broadcast at the sensing-server's frame
    cadence. Persons are a list of opaque dicts — typed PoseEstimate
    decoding lives in the P2 bindings; the WS client passes through."""
    node_id: str = ""
    timestamp: float = 0.0
    persons: tuple[dict[str, Any], ...] = ()
    confidence: float = 0.0


# ─── Decoder ─────────────────────────────────────────────────────────


def _decode(raw_text: str) -> SensingMessage:
    """Decode a single WS frame into a typed message.

    Unknown ``type`` values yield a plain ``SensingMessage`` rather
    than raising — the sensing-server is on a faster release cadence
    than this client, and unknown types should not break the stream.
    """
    obj = json.loads(raw_text)
    if not isinstance(obj, dict):
        raise ValueError(f"sensing-server emitted non-dict payload: {type(obj).__name__}")
    mtype = obj.get("type", "")
    if mtype == "connection_established":
        return ConnectionEstablishedMessage(
            type=mtype,
            raw=obj,
            node_id=obj.get("node_id", ""),
            version=obj.get("version", ""),
            capabilities=tuple(obj.get("capabilities", ())),
        )
    if mtype == "edge_vitals":
        return EdgeVitalsMessage(
            type=mtype,
            raw=obj,
            node_id=obj.get("node_id", ""),
            presence=bool(obj.get("presence", False)),
            fall_detected=bool(obj.get("fall_detected", False)),
            motion=float(obj.get("motion", 0.0)),
            breathing_rate_bpm=(
                float(obj["breathing_rate_bpm"])
                if obj.get("breathing_rate_bpm") is not None else None
            ),
            heartrate_bpm=(
                float(obj["heartrate_bpm"])
                if obj.get("heartrate_bpm") is not None else None
            ),
            n_persons=int(obj.get("n_persons", 0)),
            motion_energy=float(obj.get("motion_energy", 0.0)),
            presence_score=float(obj.get("presence_score", 0.0)),
            rssi=(float(obj["rssi"]) if obj.get("rssi") is not None else None),
        )
    if mtype == "pose_data":
        persons = obj.get("persons", ())
        return PoseDataMessage(
            type=mtype,
            raw=obj,
            node_id=obj.get("node_id", ""),
            timestamp=float(obj.get("timestamp", 0.0)),
            persons=tuple(persons) if isinstance(persons, list) else (),
            confidence=float(obj.get("confidence", 0.0)),
        )
    return SensingMessage(type=mtype, raw=obj)


# ─── Client ──────────────────────────────────────────────────────────


class SensingClient:
    """Asyncio WebSocket client for the RuView sensing-server.

    Usage as async context manager:

    ```python
    async with SensingClient("ws://localhost:8765/ws/sensing") as c:
        async for msg in c.stream():
            ...
    ```

    The client does NOT auto-reconnect — if you want resilience, wrap
    the ``async with`` in your own retry loop. Auto-reconnect logic is
    application-specific (e.g., "retry forever" for a long-running
    automation vs "fail fast" for a CLI tool that should exit).
    """

    def __init__(
        self,
        url: str,
        *,
        ping_interval: float = 20.0,
        ping_timeout: float = 20.0,
        max_size: int = 16 * 1024 * 1024,
    ) -> None:
        if not _WEBSOCKETS_AVAILABLE:
            raise ImportError(
                "SensingClient requires the `websockets` package. Install with "
                "`pip install \"wifi-densepose[client]\"` to enable the client extras."
            )
        self.url = url
        self._ping_interval = ping_interval
        self._ping_timeout = ping_timeout
        self._max_size = max_size
        self._ws: Any = None  # websockets.WebSocketClientProtocol — typed Any to avoid import cost

    async def __aenter__(self) -> "SensingClient":
        self._ws = await websockets.connect(
            self.url,
            ping_interval=self._ping_interval,
            ping_timeout=self._ping_timeout,
            max_size=self._max_size,
        )
        return self

    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
        await self.close()

    async def close(self) -> None:
        """Idempotent connection close."""
        if self._ws is not None:
            try:
                await self._ws.close()
            except Exception as e:  # pragma: no cover — best-effort close
                log.debug("ignored WS close error: %r", e)
            self._ws = None

    async def stream(self) -> AsyncIterator[SensingMessage]:
        """Yield typed messages until the server closes the connection
        or the context is exited.

        Decode failures on individual frames are logged at WARN and
        swallowed — a malformed frame should not terminate the stream
        (the next frame may be fine)."""
        if self._ws is None:
            raise RuntimeError("SensingClient not connected. Use `async with` first.")
        try:
            async for frame in self._ws:
                if isinstance(frame, bytes):
                    frame = frame.decode("utf-8", errors="replace")
                try:
                    yield _decode(frame)
                except (ValueError, json.JSONDecodeError) as e:
                    log.warning("dropping malformed sensing-server frame: %r", e)
        except ConnectionClosed:
            # Graceful EOF — exit the iterator normally.
            return

    async def send_ping(self) -> None:
        """Send an application-level ping. The sensing-server replies
        with `{"type": "pong"}` (main.rs:2698)."""
        if self._ws is None:
            raise RuntimeError("SensingClient not connected. Use `async with` first.")
        await self._ws.send(json.dumps({"type": "ping"}))

    async def recv_one(self, *, timeout: Optional[float] = None) -> SensingMessage:
        """Receive a single decoded message. Convenience for short
        scripts and tests that don't need an async generator."""
        if self._ws is None:
            raise RuntimeError("SensingClient not connected. Use `async with` first.")
        if timeout is None:
            frame = await self._ws.recv()
        else:
            frame = await asyncio.wait_for(self._ws.recv(), timeout=timeout)
        if isinstance(frame, bytes):
            frame = frame.decode("utf-8", errors="replace")
        return _decode(frame)