"""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 get_container_status from src.services.tunnel import check_tunnel_health from src.services.event_bus import InstanceEventBus, InstanceEventPayload from src.services.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.""" 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) # Create notification for instance owner (fire-and-forget) # Only send warnings and errors; skip "recovered" info notifications. if new_status == "error": 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}, )