headquarter/apps/api/src/services/instance/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 import HealthCheck
from src.models import ToolInstance
from src.services.shared.correlation import get_correlation_id
from src.services.docker import get_container_status
from src.services.shared.tunnel import check_tunnel_health
from src.services.instance.event_bus import InstanceEventBus, InstanceEventPayload
from src.services.shared.notification_service import notification_service

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."""
        # Skip instances that have never been assigned a container.
        if not instance.container_id:
            logger.debug(
                "Skipping health check for instance %s: no container_id",
                instance.id,
            )
            return

        try:
            container_info = get_container_status(instance.container_id)
        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,
            previous,
            instance.status,
        )

        # If first check or snapshot changed
        if previous is None:
            # Monitor restart / first time seeing this instance.
            # Only act if the derived status is different from the DB status.
            # This prevents duplicate notifications after monitor restarts.
            if new_status == instance.status:
                self._last_known_state[instance.id] = snapshot
                return
        elif self._snapshots_equal(previous, snapshot):
            # Nothing changed since last poll — skip entirely.
            return

        await self._handle_state_change(
            session, instance, previous, snapshot, new_status
        )
        self._last_known_state[instance.id] = snapshot

    def _derive_status(
        self,
        snapshot: HealthSnapshot,
        previous: HealthSnapshot | None,
        current_status: str | None,
    ) -> str:
        """Derive instance status from health snapshot.

        Treats missing containers as an error only when the container was
        previously known to be running. This avoids false "container failed"
        alerts for instances that are still starting or have no container yet.
        """
        if snapshot.container_status == "not_found":
            # Still starting — container may not exist yet.
            if current_status == "starting":
                return "starting"
            # Container disappeared while it was supposed to be running.
            if current_status == "running":
                return "error"
            # If we have previous memory and the container was running,
            # mark as error (handles monitor restart edge case).
            if previous is not None and previous.container_status == "running":
                return "error"
            # Fall back to current status to avoid spurious errors.
            return current_status or "error"

        if snapshot.container_status in ("exited", "dead"):
            return "error"

        if snapshot.container_status == "running":
            if snapshot.tunnel_healthy is False:
                return "unhealthy"
            return "running"

        # Transient states (created, restarting) — preserve current status
        # instead of treating them as an error. The next poll will resolve.
        if snapshot.container_status in ("created", "restarting"):
            # A container that was already running and is now restarting has
            # crashed (e.g. entrypoint failure / restart loop). Mark it failed
            # so the dashboard does not keep showing it as running.
            if current_status == "running":
                return "error"
            return current_status or "starting"

        # Unknown/unexpected state (paused, etc.)
        return "error"

    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"
            if new_status == "unhealthy":
                message = "Container tunnel is unreachable"
            else:
                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)

        # Create notification for instance owner (fire-and-forget)
        # Only send warnings and errors; skip "recovered" info notifications.
        if new_status == "error":
            # Skip duplicate error notifications if already in error state.
            if previous_status == "error":
                return
            # Skip "not_found" errors for containers that were never running
            # (e.g. still starting, or intentionally stopped/deleted).
            if snapshot.container_status == "not_found" and (
                previous is None or previous.container_status != "running"
            ):
                return
            category = "instance"
            severity = "error"
            title = "Container failed"
        elif new_status == "unhealthy":
            category = "health"
            severity = "warning"
            title = "Container unhealthy"
        else:
            # Running/recovered — do not notify
            return

        try:
            await notification_service.create_notification(
                session=session,
                user_id=instance.owner_id,
                category=category,
                severity=severity,
                title=title,
                message=message,
                source_type="tool_instances",
                source_id=instance.id,
                metadata=metadata,
            )
        except Exception:
            logger.exception(
                "Failed to create notification for health event %s",
                event_type,
                extra={"correlation_id": correlation_id},
            )