rein-aharness/rein_aharness/hub.py

"""Custodian State Hub reporting (REST, no MCP).

STATE_HUB_URL defaults to the workstation-local hub; on Railiance the
ops-bridge tunnel exposes it at http://127.0.0.1:18000.
"""

from __future__ import annotations

import os
from typing import Any

import httpx

_DEFAULT_URL = "http://127.0.0.1:8000"
_TIMEOUT = 10.0


def _base_url() -> str:
    return os.environ.get("STATE_HUB_URL", _DEFAULT_URL).rstrip("/")


def post_progress_event(
    summary: str,
    event_type: str,
    detail: dict[str, Any],
    task_id: str | None = None,
) -> bool:
    payload: dict[str, Any] = {
        "summary": summary,
        "event_type": event_type,
        "detail": detail,
        "author": "agt-executor-worker",
    }
    if task_id:
        payload["task_id"] = task_id
    try:
        resp = httpx.post(f"{_base_url()}/progress/", json=payload, timeout=_TIMEOUT)
        resp.raise_for_status()
        return True
    except httpx.HTTPError:
        return False


def close_task(task_id: str) -> bool:
    try:
        resp = httpx.patch(
            f"{_base_url()}/tasks/{task_id}",
            json={"status": "done"},
            timeout=_TIMEOUT,
        )
        resp.raise_for_status()
        return True
    except httpx.HTTPError:
        return False


def post_token_event(
    repo: str,
    tokens: int,
    *,
    budget: int | None = None,
    agent: str | None = None,
    ok: bool = True,
    detail: dict[str, Any] | None = None,
) -> bool:
    """Best-effort token cost event for the hub Token Cost dashboard.

    Schema is intentionally loose: if the hub rejects the payload we
    swallow the error so a metrics-schema drift never fails a run.
    """
    payload: dict[str, Any] = {
        "repo": repo,
        "tokens": tokens,
        "source": "rein-aharness",
        "ok": ok,
    }
    if budget is not None:
        payload["budget"] = budget
    if agent is not None:
        payload["agent"] = agent
    if detail:
        payload["detail"] = detail
    try:
        resp = httpx.post(
            f"{_base_url()}/token-events/upsert",
            json=payload,
            timeout=_TIMEOUT,
        )
        if resp.status_code >= 400:
            # Fallback shape used by some hub builds
            resp = httpx.post(
                f"{_base_url()}/token-events/",
                json=payload,
                timeout=_TIMEOUT,
            )
        resp.raise_for_status()
        return True
    except httpx.HTTPError:
        return False