LetsBeBiz-Redesign/letsbe-sysadmin-agent/app/clients/hub_client.py

"""Async HTTP client for communicating with the LetsBe Hub."""

import asyncio
from typing import Any, Optional

import httpx

from app.config import Settings, get_settings
from app.utils.credential_reader import get_all_tool_credentials, get_credential_hash
from app.utils.logger import get_logger

logger = get_logger("hub_client")


class HubClient:
    """Async client for Hub REST API.

    Used for sending heartbeats with tool credentials directly to the Hub.
    This bypasses the orchestrator for credential synchronization.
    """

    def __init__(self, settings: Optional[Settings] = None):
        self.settings = settings or get_settings()
        self._client: Optional[httpx.AsyncClient] = None
        self._last_credentials_hash: str = ""

    @property
    def is_configured(self) -> bool:
        """Check if Hub connection is configured."""
        return bool(
            self.settings.hub_url
            and self.settings.hub_api_key
            and self.settings.hub_telemetry_enabled
        )

    def _get_headers(self) -> dict[str, str]:
        """Get headers for Hub API requests."""
        return {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.settings.hub_api_key}",
            "X-Agent-Version": self.settings.agent_version,
            "X-Agent-Hostname": self.settings.hostname,
        }

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create the HTTP client."""
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self.settings.hub_url,
                headers=self._get_headers(),
                timeout=httpx.Timeout(30.0, connect=10.0),
            )
        return self._client

    async def send_heartbeat(
        self,
        include_credentials: bool = True,
        status: Optional[dict[str, Any]] = None,
    ) -> bool:
        """Send heartbeat to Hub with optional credentials.

        Args:
            include_credentials: Include tool credentials in heartbeat
            status: Optional system status metrics

        Returns:
            True if heartbeat was sent successfully
        """
        if not self.is_configured:
            logger.debug("hub_heartbeat_skipped", reason="not_configured")
            return False

        try:
            payload: dict[str, Any] = {
                "agentVersion": self.settings.agent_version,
            }

            # Include system status if provided
            if status:
                payload["status"] = status

            # Include tool credentials only when they've changed
            if include_credentials:
                current_hash = get_credential_hash()
                if current_hash and current_hash != self._last_credentials_hash:
                    credentials = get_all_tool_credentials()
                    if credentials:
                        payload["credentials"] = credentials
                        payload["credentialsHash"] = current_hash
                        self._last_credentials_hash = current_hash
                        logger.debug(
                            "hub_heartbeat_with_credentials",
                            tools=list(credentials.keys()),
                        )
                elif current_hash:
                    # Just send the hash so Hub knows credentials haven't changed
                    payload["credentialsHash"] = current_hash

            client = await self._get_client()
            response = await client.post(
                "/api/v1/orchestrator/heartbeat",
                json=payload,
            )

            if response.status_code == 200:
                data = response.json()
                logger.info(
                    "hub_heartbeat_sent",
                    server_id=data.get("serverId"),
                    commands_pending=len(data.get("commands", [])),
                )
                return True
            elif response.status_code == 401:
                logger.warning(
                    "hub_heartbeat_auth_failed",
                    status_code=response.status_code,
                )
                return False
            else:
                logger.warning(
                    "hub_heartbeat_failed",
                    status_code=response.status_code,
                    response=response.text[:200],
                )
                return False

        except (httpx.ConnectError, httpx.TimeoutException) as e:
            logger.warning("hub_heartbeat_network_error", error=str(e))
            return False
        except Exception as e:
            logger.error("hub_heartbeat_error", error=str(e))
            return False

    async def close(self) -> None:
        """Close the HTTP client."""
        if self._client and not self._client.is_closed:
            await self._client.aclose()
            self._client = None


# Singleton instance
_hub_client: Optional[HubClient] = None


def get_hub_client() -> HubClient:
    """Get the singleton Hub client instance."""
    global _hub_client
    if _hub_client is None:
        _hub_client = HubClient()
    return _hub_client


async def send_hub_heartbeat() -> bool:
    """Convenience function to send heartbeat to Hub.

    Returns:
        True if heartbeat was sent successfully, False if not configured or failed
    """
    client = get_hub_client()
    return await client.send_heartbeat()