headquarter/apps/api/src/services/health_monitor.py

"""Background health monitor that polls container and tunnel health."""

import asyncio
import logging
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from src.database import SessionLocal
from src.models.health_check import HealthCheck
from src.models.tool_instance import ToolInstance
from src.services.correlation import get_correlation_id
from src.services.docker import check_tunnel_health, get_container_status
from src.services.event_bus import InstanceEventBus, InstanceEventPayload

logger = logging.getLogger(__name__)


@dataclass
class HealthSnapshot:
    """In-memory snapshot of an instance's health state."""

    container_status: str | None = None
    container_healthy: bool | None = None
    tunnel_healthy: bool | None = None
    exit_code: int | None = None


class HealthMonitor:
    """Polls container and tunnel health, publishing events on state changes."""

    POLL_INTERVAL_SECONDS: float = 15.0
    _MONITORED_STATUSES: set[str] = {"starting", "running", "unhealthy"}

    def __init__(self, event_bus: InstanceEventBus) -> None:
        self._event_bus = event_bus
        self._task: asyncio.Task | None = None
        self._last_known_state: dict[uuid.UUID, HealthSnapshot] = {}

    def start(self) -> None:
        """Idempotent start of the background polling task."""
        if self._task is not None and not self._task.done():
            return
        try:
            loop = asyncio.get_running_loop()
            self._task = loop.create_task(self._poll_loop())
        except RuntimeError:
            pass

    def stop(self) -> None:
        """Cancel the background task and clear state."""
        if self._task is not None and not self._task.done():
            self._task.cancel()
        self._last_known_state.clear()
        self._task = None

    async def _poll_loop(self) -> None:
        """Main polling loop."""
        while True:
            try:
                await asyncio.sleep(self.POLL_INTERVAL_SECONDS)
                await self._run_check_cycle()
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("Health monitor poll loop error")

    async def _run_check_cycle(self) -> None:
        """Check all monitored instances in one cycle."""
        async with SessionLocal() as session:
            result = await session.execute(
                select(ToolInstance).where(
                    ToolInstance.status.in_(self._MONITORED_STATUSES)
                )
            )
            instances = result.scalars().all()

        for instance in instances:
            async with SessionLocal() as session:
                await self._check_instance(session, instance)

    async def _check_instance(
        self,
        session: AsyncSession,
        instance: ToolInstance,
    ) -> None:
        """Check a single instance and handle state transitions."""
        try:
            container_info = get_container_status(instance.container_id or "")
        except Exception:
            logger.exception(
                "Health check failed for instance %s",
                instance.id,
                extra={
                    "instance_id": str(instance.id),
                    "correlation_id": get_correlation_id(),
                },
            )
            return

        container_status = container_info["status"]
        exit_code = container_info["exit_code"]
        container_healthy = (
            container_info["health"] == "healthy" if container_info["health"] else None
        )

        tunnel_healthy: bool | None = None
        if instance.public_url and container_status == "running":
            try:
                tunnel_result = check_tunnel_health(instance.public_url)
                tunnel_healthy = tunnel_result.get("healthy", False)
            except Exception:
                logger.exception(
                    "Tunnel health check failed for instance %s",
                    instance.id,
                    extra={
                        "instance_id": str(instance.id),
                        "correlation_id": get_correlation_id(),
                    },
                )
                tunnel_healthy = False

        snapshot = HealthSnapshot(
            container_status=container_status,
            container_healthy=container_healthy,
            tunnel_healthy=tunnel_healthy,
            exit_code=exit_code,
        )

        previous = self._last_known_state.get(instance.id)

        # Determine new status
        new_status = self._derive_status(snapshot)

        # If first check or state changed
        if previous is None or not self._snapshots_equal(previous, snapshot):
            await self._handle_state_change(
                session, instance, previous, snapshot, new_status
            )
            self._last_known_state[instance.id] = snapshot

    def _derive_status(self, snapshot: HealthSnapshot) -> str:
        """Derive instance status from health snapshot."""
        if snapshot.container_status != "running":
            return "error"
        if snapshot.tunnel_healthy is False:
            return "unhealthy"
        return "running"

    def _snapshots_equal(self, a: HealthSnapshot, b: HealthSnapshot) -> bool:
        """Compare two snapshots for equality."""
        return (
            a.container_status == b.container_status
            and a.container_healthy == b.container_healthy
            and a.tunnel_healthy == b.tunnel_healthy
            and a.exit_code == b.exit_code
        )

    async def _handle_state_change(
        self,
        session: AsyncSession,
        instance: ToolInstance,
        previous: HealthSnapshot | None,
        snapshot: HealthSnapshot,
        new_status: str,
    ) -> None:
        """Update DB, insert health check, and publish event."""
        previous_status = instance.status

        # Update instance status
        instance.status = new_status
        if new_status == "error":
            instance.last_stopped_at = datetime.now(timezone.utc)

        # Insert health check row
        health_check = HealthCheck(
            instance_id=instance.id,
            container_status=snapshot.container_status,
            container_healthy=snapshot.container_healthy,
            tunnel_healthy=snapshot.tunnel_healthy,
            exit_code=snapshot.exit_code,
            probe_status=None,
            probe_output=None,
        )
        session.add(health_check)
        await session.commit()

        # Build event payload
        correlation_id = get_correlation_id()
        metadata: dict = {"previous_status": previous_status}
        if snapshot.exit_code is not None:
            metadata["exit_code"] = snapshot.exit_code
            metadata["error_type"] = "container"
        if instance.public_url:
            metadata["tunnel_url"] = instance.public_url

        if new_status == "error":
            event_type = "instance.error"
            message = f"Container failed with status {snapshot.container_status}"
            if snapshot.exit_code is not None:
                message += f" (exit code: {snapshot.exit_code})"
        else:
            event_type = "instance.health_changed"
            message = f"Container is now {new_status}"

        payload: InstanceEventPayload = {
            "event": event_type,
            "instance_id": str(instance.id),
            "status": new_status,
            "message": message,
            "metadata": metadata,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "correlation_id": correlation_id,
        }

        await self._event_bus.publish(event_type, payload)