161 lines
5.3 KiB
Python
161 lines
5.3 KiB
Python
"""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()
|