Source code for ralph.pro_support.heartbeat

"""Bounded Pro heartbeat client.

The engine, when running as a Pro subprocess, POSTs a small JSON
heartbeat to ``<base_url>/api/heartbeat`` every ``interval_seconds``
seconds so Pro can monitor liveness. The heartbeat client is a
self-contained class that:

- runs the heartbeat loop in a **daemon thread** so the process can
  always exit even if Pro is hung;
- uses an **explicit bounded ``timeout=``** on every ``httpx`` call so
  the bounded-subprocess audit (``ralph.testing.audit_mcp_timeout``)
  catches any regression;
- treats ``401`` and ``404`` responses as **hard stops** — once the
  heartbeat is rejected as unauthorized or unknown, the client logs a
  warning and stops looping;
- treats every other error (connection refused, timeout, 5xx) as
  **transient** — log at debug level and continue, so a Pro restart or
  brief outage does not crash the pipeline;
- exposes an **idempotent ``stop()``** that only sets a
  ``threading.Event``; it does NOT join the worker because daemon
  threads cannot be meaningfully joined and the process must never
  block on a slow Pro server.

The client does not perform I/O at construction time. ``start()``
launches the daemon thread. ``stop()`` is safe to call multiple times
and from any thread.
"""

from __future__ import annotations

import logging
import threading
import time
from typing import TYPE_CHECKING, Protocol

if TYPE_CHECKING:
    from collections.abc import Callable, Mapping


logger = logging.getLogger(__name__)


# HTTP status codes for the Pro heartbeat endpoint. Defined as
# module-level constants so the magic-number check in
# ``_post_once`` does not trigger PLR2004.
_STATUS_UNAUTHORIZED = 401
_STATUS_NOT_FOUND = 404
_STATUS_SERVER_ERROR_THRESHOLD = 500


class _HttpxResponseLike(Protocol):
    """Minimum surface we need from a httpx-like response object."""

    status_code: int


class _HttpxClientLike(Protocol):
    """Minimum surface we need from a httpx-like client."""

    def post(
        self,
        url: str,
        *,
        json: dict[str, object] | None = None,
        timeout: float | None = None,
    ) -> _HttpxResponseLike: ...


class _HttpxClientFactory(Protocol):
    def __call__(self) -> _HttpxClientLike: ...


def _default_httpx_client_factory() -> _HttpxClientLike:
    """Return a default ``httpx.Client`` for production use.

    The factory indirection exists solely so tests can substitute a
    fake without monkeypatching the ``httpx`` module. The import is
    deliberately lazy so the pro_support module remains importable
    even when ``httpx`` is not installed (e.g. during lightweight
    test runs).
    """
    import httpx  # noqa: PLC0415 - lazy import to keep pro_support importable without httpx

    # Bound the default client at construction so a slow Pro server
    # can never block the daemon thread at the connection layer; the
    # per-request timeout is still applied via ``timeout=self._timeout``
    # in ``_post_once``.
    return httpx.Client(timeout=5.0)  # resource-lifecycle-ok: closed in _post_once finally


def _default_clock() -> float:
    return time.monotonic()


[docs] class ProHeartbeatClient: """Bounded, daemon-threaded heartbeat client for the Pro subprocess contract. Constructor parameters are explicit (no module-level mutable state) so every test can construct a client with a fake clock and a fake httpx factory. """ def __init__( self, run_id: str, token: str, base_url: str, pid: int, *, interval_seconds: float = 5.0, timeout_seconds: float = 5.0, httpx_client_factory: _HttpxClientFactory | None = None, clock: Callable[[], float] | None = None, metadata: Mapping[str, object] | None = None, ) -> None: if interval_seconds <= 0: raise ValueError("interval_seconds must be positive") if timeout_seconds <= 0: raise ValueError("timeout_seconds must be positive") self._run_id = run_id self._token = token self._base_url = base_url.rstrip("/") self._pid = pid self._interval = float(interval_seconds) self._timeout = float(timeout_seconds) self._client_factory: _HttpxClientFactory = ( httpx_client_factory if httpx_client_factory is not None else _default_httpx_client_factory ) self._clock: Callable[[], float] = clock if clock is not None else _default_clock self._metadata: dict[str, object] = dict(metadata) if metadata is not None else {} self._stop_event = threading.Event() self._thread: threading.Thread | None = None
[docs] def start(self) -> None: """Spawn the daemon worker thread. Idempotent: a second call is a no-op.""" if self._thread is not None and self._thread.is_alive(): return self._stop_event.clear() thread = threading.Thread( target=self._run_loop, name="ralph-pro-heartbeat", daemon=True, ) self._thread = thread thread.start()
[docs] def stop(self) -> None: """Signal the worker to exit on its next loop iteration. Idempotent. Deliberately does NOT ``join()``: the worker is daemonic and therefore will be torn down with the process; joining a daemon thread can block on a slow Pro server which would defeat the entire point of the design. """ self._stop_event.set()
@property def is_running(self) -> bool: return self._thread is not None and self._thread.is_alive() def _build_payload(self) -> dict[str, object]: return { "run_id": self._run_id, "token": self._token, "status": "running", "pid": self._pid, "metadata": dict(self._metadata), } def _run_loop(self) -> None: next_tick_at = self._clock() while not self._stop_event.is_set(): now = self._clock() if now < next_tick_at: self._sleep_for(next_tick_at - now) continue self._post_once() next_tick_at = self._clock() + self._interval def _sleep_for(self, seconds: float) -> None: """Sleep until the stop event fires or ``seconds`` elapses. Implemented with ``Event.wait(timeout=...)`` so a ``stop()`` call from the main thread can interrupt the sleep immediately rather than waiting out the full interval. ``Event.wait`` is bounded by its timeout so it cannot wedge the daemon. """ self._stop_event.wait(timeout=seconds) def _post_once(self) -> None: url = f"{self._base_url}/api/heartbeat" payload = self._build_payload() try: client = self._client_factory() except Exception as exc: logger.debug("Pro heartbeat client creation failed (transient): %s", exc) return try: try: response = client.post(url, json=payload, timeout=self._timeout) except Exception as exc: logger.debug("Pro heartbeat POST to %s failed (transient): %s", url, exc) return status_code: int = response.status_code if status_code in (_STATUS_UNAUTHORIZED, _STATUS_NOT_FOUND): logger.warning( "Pro heartbeat rejected with status %s; stopping heartbeat loop", status_code, ) self._stop_event.set() return if status_code >= _STATUS_SERVER_ERROR_THRESHOLD: logger.debug( "Pro heartbeat transient server error %s; continuing", status_code, ) return finally: close_method: Callable[[], object] | None = getattr(client, "close", None) if callable(close_method): try: close_method() except Exception: logger.debug("Pro heartbeat client close failed", exc_info=True)
__all__ = ["ProHeartbeatClient"]