cloudagent-daemon 0.2.1__py3-none-any.whl

cloudagent_daemon/ws_client.py

@@ -0,0 +1,431 @@
+"""Async WebSocket client for the cloudagent daemon.
+
+Spec: docs/superpowers/specs/2026-05-08-local-daemon-runtime-design.md §4.1
+
+Single responsibility: own the WS transport. Connect → hello → ready
+handshake, run a background reader that feeds an inbound queue, auto-reply
+to server pings, and watchdog the connection (>40s of silence ⇒ close).
+
+Also exposes ``run_forever`` — a reconnect-on-disconnect supervisor with
+exponential backoff (1→2→4→8→16→30s cap) that sends ``resume {last_seq}``
+after every reconnect so the server can replay unacked frames (Task 5).
+
+Out of scope (handled in later tasks):
+  - Frame dispatch by type (Task 6)
+  - Outbound queue + delivery ack (Task 7)
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import platform
+import socket
+import ssl
+import time
+from collections.abc import Awaitable, Callable
+from datetime import UTC, datetime
+from typing import Any
+from uuid import uuid4
+
+from websockets.asyncio.client import ClientConnection, connect
+from websockets.exceptions import ConnectionClosed, InvalidStatus
+
+from .runtime_detect import DetectedRuntime, detect_runtimes
+
+log = logging.getLogger(__name__)
+
+# Spec §4.1: server pings every 30s. We treat 40s of silence as "missed
+# two pings" and close. Watchdog fires every 5s.
+HEARTBEAT_GRACE_SECONDS = 40.0
+WATCHDOG_INTERVAL_SECONDS = 5.0
+
+
+def _ws_url(server_url: str) -> str:
+    """Map ``http(s)://host`` → ``ws(s)://host/v1/daemon/ws``."""
+    base = server_url.rstrip("/")
+    if base.startswith("https://"):
+        base = "wss://" + base[len("https://"):]
+    elif base.startswith("http://"):
+        base = "ws://" + base[len("http://"):]
+    return f"{base}/v1/daemon/ws"
+
+
+def _new_frame_id() -> str:
+    return f"frm_{uuid4().hex[:16]}"
+
+
+def _now_iso() -> str:
+    return datetime.now(UTC).isoformat()
+
+
+class WSClient:
+    """Asyncio WebSocket client for ``/v1/daemon/ws``.
+
+    Use as ``async with WSClient(server_url, api_key) as client:``. The
+    handshake (hello + ready) runs inside ``__aenter__``; on exit we
+    cancel the reader/watchdog tasks and close the WS.
+    """
+
+    def __init__(
+        self,
+        server_url: str,
+        api_key: str,
+        *,
+        daemon_ver: str | None = None,
+        detect_runtimes_fn: Callable[[], list[DetectedRuntime]] | None = None,
+        insecure: bool = False,
+        workspace_root: str | None = None,
+    ) -> None:
+        # Default the version to the package's __version__ so the WS hello
+        # frame stays in sync with the SSE/POST ready POST without
+        # callers having to remember to pass it.
+        if daemon_ver is None:
+            from cloudagent_daemon import __version__ as _v
+            daemon_ver = _v
+        self._server_url = server_url
+        self._api_key = api_key
+        self._daemon_ver = daemon_ver
+        # Task #113 (2026-05-12) — agent workspace root, surfaced on
+        # the legacy WS ``ready`` frame so the legacy path mirrors the
+        # SSE+POST ``/v1/daemon/ready`` shape. ``None`` ⇒ omit the
+        # field (backend preserves any previously-reported value).
+        self._workspace_root = workspace_root
+        # DI seam — tests pass a stub returning a known list so the ready
+        # frame is deterministic without monkeypatching the module. Defaults
+        # to scanning $PATH on the real machine.
+        self._detect_runtimes_fn = detect_runtimes_fn or detect_runtimes
+        # TLS verification policy. Default = strict (verify cert + hostname).
+        # When ``insecure=True`` we build an unverified SSL context — required
+        # for self-signed dev/staging environments (e.g. fr cluster pre-Wave-7
+        # where nginx-ingress serves the bundled fake cert). Only applies to
+        # wss:// URLs; for ws:// the ssl context is unused.
+        self._insecure = insecure
+        if insecure:
+            self._ssl_context: ssl.SSLContext | None = ssl.create_default_context()
+            self._ssl_context.check_hostname = False
+            self._ssl_context.verify_mode = ssl.CERT_NONE
+        else:
+            # Let websockets build the default verified context per scheme.
+            self._ssl_context = None
+        self._ws: ClientConnection | None = None
+        self._inbound: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+        self._reader_task: asyncio.Task[None] | None = None
+        self._watchdog_task: asyncio.Task[None] | None = None
+        self._host_id: str = ""
+        self._bound_agents: list[dict[str, Any]] = []
+        self._last_inbound_at: float = 0.0
+        self._closed = False
+        # Set when the reader loop has exited (server dropped, etc.) so
+        # next_frame() can surface the disconnect instead of hanging.
+        self._disconnect_event: asyncio.Event = asyncio.Event()
+
+    @property
+    def host_id(self) -> str:
+        return self._host_id
+
+    @property
+    def bound_agents(self) -> list[dict[str, Any]]:
+        return list(self._bound_agents)
+
+    def _build_ready_frame(
+        self, runtimes_payload: list[dict[str, Any]]
+    ) -> dict[str, Any]:
+        """Build the ``ready`` frame body sent during handshake.
+
+        Pulled out as a helper so tests can verify the payload shape
+        (notably Task #113's ``workspace_root`` field) without driving
+        the full websocket handshake.
+        """
+        ready_payload: dict[str, Any] = {
+            "daemon_ver": self._daemon_ver,
+            "runtimes": runtimes_payload,
+            "hostname": socket.gethostname(),
+            "os": platform.system().lower(),
+        }
+        # Task #113 — only emit when set; omission preserves backend
+        # state (idempotent re-ready).
+        if self._workspace_root is not None:
+            ready_payload["workspace_root"] = self._workspace_root
+        return {
+            "v": 1,
+            "type": "ready",
+            "id": _new_frame_id(),
+            "ts": _now_iso(),
+            "payload": ready_payload,
+        }
+
+    async def __aenter__(self) -> WSClient:
+        # ``additional_headers`` is the websockets v15 spelling (renamed
+        # from ``extra_headers``). ``ping_interval=None`` disables the
+        # library's auto-ping; we implement application-level heartbeat
+        # ourselves per spec §4.1.
+        ws_url = _ws_url(self._server_url)
+        # Only pass ssl= for wss://. For plain ws:// websockets ignores it
+        # but is fussy about None vs default in some library versions.
+        connect_kwargs: dict[str, Any] = {
+            "additional_headers": {"Authorization": f"Bearer {self._api_key}"},
+            "ping_interval": None,
+        }
+        if ws_url.startswith("wss://") and self._ssl_context is not None:
+            connect_kwargs["ssl"] = self._ssl_context
+        self._ws = await connect(ws_url, **connect_kwargs)
+
+        # Hello is the first inbound frame from the server.
+        try:
+            raw = await self._ws.recv()
+        except ConnectionClosed as exc:
+            await self._ws.close()
+            raise ConnectionError(
+                f"server closed before hello: code={exc.rcvd}"
+            ) from exc
+        hello = json.loads(raw)
+        if hello.get("type") != "hello":
+            await self._ws.close()
+            raise ConnectionError(
+                f"expected hello frame, got {hello.get('type')!r}"
+            )
+        payload = hello.get("payload", {}) or {}
+        self._host_id = payload.get("host_id", "")
+        self._bound_agents = list(payload.get("agents", []) or [])
+        self._last_inbound_at = time.monotonic()
+
+        # Detect locally installed AI coding runtimes (Task 6). Failures
+        # in detection (e.g. a probe binary hangs unexpectedly) must not
+        # break the handshake — degrade to empty list so the daemon still
+        # comes online and the user can fix runtime config out-of-band.
+        try:
+            detected = self._detect_runtimes_fn()
+        except Exception as exc:  # pragma: no cover — defensive
+            log.warning("runtime detection failed (%s); reporting empty list", exc)
+            detected = []
+        runtimes_payload = [r.to_payload() for r in detected]
+
+        ready_frame = self._build_ready_frame(runtimes_payload)
+        await self._ws.send(json.dumps(ready_frame))
+
+        self._reader_task = asyncio.create_task(self._reader_loop())
+        self._watchdog_task = asyncio.create_task(self._watchdog_loop())
+        return self
+
+    async def __aexit__(self, *_exc: Any) -> None:
+        self._closed = True
+        for task in (self._reader_task, self._watchdog_task):
+            if task is not None and not task.done():
+                task.cancel()
+                try:
+                    await task
+                except (asyncio.CancelledError, Exception):
+                    pass
+        if self._ws is not None:
+            try:
+                await self._ws.close()
+            except Exception:
+                pass
+
+    async def next_frame(self, timeout: float | None = None) -> dict[str, Any]:
+        """Wait for the next non-control frame. Pings never surface here.
+
+        Raises ``ConnectionClosed`` if the underlying reader has already
+        exited (server dropped the connection) — otherwise the call would
+        hang forever on an empty queue. Callers like ``run_forever``
+        rely on this to trigger a reconnect.
+        """
+        get_task = asyncio.create_task(self._inbound.get())
+        disc_task = asyncio.create_task(self._disconnect_event.wait())
+        try:
+            done, _ = await asyncio.wait(
+                {get_task, disc_task},
+                timeout=timeout,
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+            if not done:
+                raise asyncio.TimeoutError
+            if get_task in done:
+                return get_task.result()
+            # Reader exited before a frame arrived → connection dead.
+            raise ConnectionClosed(rcvd=None, sent=None)
+        finally:
+            for t in (get_task, disc_task):
+                if not t.done():
+                    t.cancel()
+
+    async def send(self, frame: dict[str, Any]) -> None:
+        if self._ws is None:
+            raise RuntimeError("WSClient not connected")
+        await self._ws.send(json.dumps(frame))
+
+    async def send_resume(self, last_seq: int) -> None:
+        """Send a ``resume`` frame so the server can replay unacked frames.
+
+        Spec §4.1: after every (re)connect, the daemon tells the server
+        the highest seq it has durably processed; the server then replays
+        any frames with ``seq > last_seq`` from its catchup buffer.
+        """
+        await self.send({
+            "v": 1,
+            "type": "resume",
+            "id": _new_frame_id(),
+            "ts": _now_iso(),
+            "payload": {"last_seq": last_seq},
+        })
+
+    async def _reader_loop(self) -> None:
+        assert self._ws is not None
+        try:
+            async for raw in self._ws:
+                self._last_inbound_at = time.monotonic()
+                try:
+                    frame = json.loads(raw)
+                except (TypeError, ValueError):
+                    # Malformed frame — drop and continue. Phase 1c will
+                    # add a bye/disconnect path for protocol errors.
+                    continue
+                if frame.get("type") == "ping":
+                    # Inline pong reply, do NOT enqueue. Spec §4.1.
+                    try:
+                        await self._ws.send(json.dumps({
+                            "v": 1,
+                            "type": "pong",
+                            "id": _new_frame_id(),
+                            "ts": _now_iso(),
+                            "payload": {},
+                        }))
+                    except Exception:
+                        return
+                    continue
+                await self._inbound.put(frame)
+        except ConnectionClosed:
+            return
+        finally:
+            # Wake any next_frame() caller waiting on an empty queue.
+            self._disconnect_event.set()
+
+    async def _watchdog_loop(self) -> None:
+        """Close the WS if no inbound traffic for HEARTBEAT_GRACE_SECONDS."""
+        while True:
+            await asyncio.sleep(WATCHDOG_INTERVAL_SECONDS)
+            if self._closed or self._ws is None:
+                return
+            if (
+                time.monotonic() - self._last_inbound_at
+                > HEARTBEAT_GRACE_SECONDS
+            ):
+                try:
+                    await self._ws.close()
+                finally:
+                    return
+
+
+# ---------------------------------------------------------------------------
+# Reconnect supervisor (Task 5)
+# ---------------------------------------------------------------------------
+
+# Backoff cap from the spec — the longest delay we ever sleep between
+# reconnect attempts. The series is 1, 2, 4, 8, 16, 30, 30, 30, …
+_BACKOFF_INITIAL_SECONDS = 1.0
+_BACKOFF_CAP_SECONDS = 30.0
+
+
+class AuthFailed(Exception):
+    """Raised when the server rejects the WS upgrade with 401/403.
+
+    Auth failures are terminal — re-trying with the same key is futile
+    and could trigger server-side abuse mitigations. ``run_forever``
+    surfaces this instead of looping.
+    """
+
+    def __init__(self, status_code: int, message: str = "") -> None:
+        self.status_code = status_code
+        super().__init__(message or f"auth rejected (HTTP {status_code})")
+
+
+async def run_forever(
+    client_factory: Callable[[], WSClient],
+    on_frame: Callable[[dict[str, Any]], Awaitable[None]],
+    state_loader: Callable[[], int],
+    *,
+    on_connect: Callable[[WSClient], Awaitable[None]] | None = None,
+    _backoff_factor: float = 1.0,
+    _sleep: Callable[[float], Awaitable[None]] | None = None,
+) -> None:
+    """Run the WS client forever, reconnecting with exponential backoff.
+
+    Args:
+        client_factory: Builds a fresh ``WSClient`` per attempt — the
+            client is single-use because ``__aexit__`` cancels its
+            background tasks.
+        on_frame: Coroutine called for every business frame surfaced by
+            ``WSClient.next_frame``. Frame dispatch logic lives here
+            (Task 6 wires it up to the assign/deliver/tool_call handlers).
+        state_loader: Returns the current ``last_seq`` to send in the
+            ``resume`` frame after each (re)connect. Called fresh each
+            iteration so it always reflects the latest persisted state.
+        on_connect: Optional coroutine called once per successful
+            (re)connect, after the WS handshake completes and *before*
+            the resume frame is sent. Used by ``main_loop.run_daemon``
+            to rebind ``OutboundQueue.ws`` to the freshly minted client
+            so subsequent outbound sends target the live socket.
+        _backoff_factor: Multiplier applied to all sleep durations. Real
+            usage leaves this at 1.0 (delays in real seconds); tests
+            pass a small value (e.g. 0.01) to keep the suite fast.
+        _sleep: Optional sleep override — defaults to ``asyncio.sleep``.
+            Test-only hook used to record/replace backoff sleeps without
+            patching the global ``asyncio.sleep`` (which would also break
+            unrelated waits inside the same test).
+
+    Raises:
+        AuthFailed: Server returned HTTP 401/403 — auth is dead, do not
+            retry. Caller should re-prompt the user for their machine
+            key (or just exit).
+        asyncio.CancelledError: Propagated unchanged so cooperative
+            shutdown (signal handlers, ``task.cancel()``) works.
+    """
+    sleep = _sleep if _sleep is not None else asyncio.sleep
+    backoff = _BACKOFF_INITIAL_SECONDS
+    while True:
+        try:
+            client = client_factory()
+            async with client:
+                # Connection succeeded → reset backoff so the next drop
+                # starts the curve over at 1s.
+                backoff = _BACKOFF_INITIAL_SECONDS
+
+                # Notify caller (e.g. main_loop) so it can rebind any
+                # client-derived state (OutboundQueue.ws). Runs BEFORE
+                # the resume frame so the resume itself uses the new
+                # client; runs AFTER aenter so .send is usable.
+                if on_connect is not None:
+                    await on_connect(client)
+
+                # Send resume BEFORE entering the receive loop so the
+                # server replays unacked frames first; on_frame then
+                # picks them up in order. last_seq=0 still gets a
+                # resume frame — the server treats that as "send me
+                # everything in the buffer".
+                last_seq = state_loader()
+                await client.send_resume(last_seq)
+
+                while True:
+                    frame = await client.next_frame()
+                    await on_frame(frame)
+        except asyncio.CancelledError:
+            # Cooperative shutdown — bubble up so the caller (signal
+            # handler, supervisor task) sees a clean exit.
+            raise
+        except InvalidStatus as exc:
+            code = exc.response.status_code
+            if code in (401, 403):
+                # Auth is dead. Re-trying won't help and risks lockout.
+                log.error("ws auth rejected (HTTP %d) — not retrying", code)
+                raise AuthFailed(code) from exc
+            # Any other handshake-level rejection (5xx, 426, …) is
+            # plausibly transient — back off and retry.
+            log.warning("ws handshake rejected (HTTP %d), backing off", code)
+        except (ConnectionClosed, OSError, asyncio.TimeoutError) as exc:
+            log.warning("ws disconnected (%s), backing off %.1fs",
+                        type(exc).__name__, backoff)
+
+        # Sleep with the current backoff, then double (capped at 30s).
+        await sleep(min(backoff, _BACKOFF_CAP_SECONDS) * _backoff_factor)
+        backoff = min(backoff * 2, _BACKOFF_CAP_SECONDS)