agent_os_kernel 3.1.0__py3-none-any.whl

atr/health.py

@@ -0,0 +1,557 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Health check functionality for ATR tools.
+
+Provides mechanisms to verify external tools and APIs are available
+before execution.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import Any, Callable, Dict, Optional, Union
+
+
+class HealthStatus(str, Enum):
+    """Health check status values."""
+
+    HEALTHY = "healthy"
+    DEGRADED = "degraded"
+    UNHEALTHY = "unhealthy"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class HealthCheckResult:
+    """Result of a health check.
+
+    Attributes:
+        status: Overall health status.
+        message: Human-readable status message.
+        latency_ms: Check latency in milliseconds.
+        timestamp: When the check was performed.
+        details: Additional details about the check.
+    """
+
+    status: HealthStatus
+    message: str = ""
+    latency_ms: Optional[float] = None
+    timestamp: datetime = field(default_factory=datetime.now)
+    details: Dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def is_healthy(self) -> bool:
+        """Check if status indicates tool is usable."""
+        return self.status in (HealthStatus.HEALTHY, HealthStatus.DEGRADED)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "status": self.status.value,
+            "message": self.message,
+            "latency_ms": self.latency_ms,
+            "timestamp": self.timestamp.isoformat(),
+            "details": self.details,
+        }
+
+
+class HealthCheck(ABC):
+    """Abstract base class for health checks.
+
+    Subclass this to create custom health checks for your tools.
+    """
+
+    @abstractmethod
+    def check(self) -> HealthCheckResult:
+        """Perform the health check synchronously.
+
+        Returns:
+            HealthCheckResult with status and details.
+        """
+        pass
+
+    async def check_async(self) -> HealthCheckResult:
+        """Perform the health check asynchronously.
+
+        Default implementation runs sync check in executor.
+        Override for true async checks.
+
+        Returns:
+            HealthCheckResult with status and details.
+        """
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, self.check)
+
+
+class HttpHealthCheck(HealthCheck):
+    """Health check that verifies an HTTP endpoint is responding.
+
+    Example:
+        >>> check = HttpHealthCheck(
+        ...     url="https://api.example.com/health",
+        ...     timeout=5.0,
+        ...     expected_status=200
+        ... )
+        >>> result = check.check()
+        >>> print(result.status)
+    """
+
+    def __init__(
+        self,
+        url: str,
+        timeout: float = 5.0,
+        expected_status: int = 200,
+        method: str = "GET",
+        headers: Optional[Dict[str, str]] = None,
+    ):
+        """Initialize HTTP health check.
+
+        Args:
+            url: The URL to check.
+            timeout: Request timeout in seconds.
+            expected_status: Expected HTTP status code.
+            method: HTTP method to use.
+            headers: Optional headers to include.
+        """
+        self.url = url
+        self.timeout = timeout
+        self.expected_status = expected_status
+        self.method = method
+        self.headers = headers or {}
+
+    def check(self) -> HealthCheckResult:
+        """Perform HTTP health check."""
+        import urllib.error
+        import urllib.request
+
+        start_time = time.perf_counter()
+
+        try:
+            request = urllib.request.Request(self.url, method=self.method, headers=self.headers)
+
+            with urllib.request.urlopen(request, timeout=self.timeout) as response:
+                latency_ms = (time.perf_counter() - start_time) * 1000
+                status_code = response.status
+
+                if status_code == self.expected_status:
+                    return HealthCheckResult(
+                        status=HealthStatus.HEALTHY,
+                        message=f"HTTP {status_code} OK",
+                        latency_ms=latency_ms,
+                        details={"url": self.url, "status_code": status_code},
+                    )
+                else:
+                    return HealthCheckResult(
+                        status=HealthStatus.DEGRADED,
+                        message=f"Unexpected status: {status_code}",
+                        latency_ms=latency_ms,
+                        details={"url": self.url, "status_code": status_code},
+                    )
+
+        except urllib.error.URLError as e:
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            return HealthCheckResult(
+                status=HealthStatus.UNHEALTHY,
+                message=f"Connection failed: {e.reason}",
+                latency_ms=latency_ms,
+                details={"url": self.url, "error": str(e)},
+            )
+        except Exception as e:
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            return HealthCheckResult(
+                status=HealthStatus.UNHEALTHY,
+                message=f"Check failed: {str(e)}",
+                latency_ms=latency_ms,
+                details={"url": self.url, "error": str(e)},
+            )
+
+
+class TcpHealthCheck(HealthCheck):
+    """Health check that verifies a TCP port is accepting connections.
+
+    Example:
+        >>> check = TcpHealthCheck(host="localhost", port=5432)
+        >>> result = check.check()
+    """
+
+    def __init__(self, host: str, port: int, timeout: float = 5.0):
+        """Initialize TCP health check.
+
+        Args:
+            host: The host to connect to.
+            port: The port to check.
+            timeout: Connection timeout in seconds.
+        """
+        self.host = host
+        self.port = port
+        self.timeout = timeout
+
+    def check(self) -> HealthCheckResult:
+        """Perform TCP health check."""
+        import socket
+
+        start_time = time.perf_counter()
+
+        try:
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.settimeout(self.timeout)
+            result = sock.connect_ex((self.host, self.port))
+            sock.close()
+
+            latency_ms = (time.perf_counter() - start_time) * 1000
+
+            if result == 0:
+                return HealthCheckResult(
+                    status=HealthStatus.HEALTHY,
+                    message=f"Port {self.port} is open",
+                    latency_ms=latency_ms,
+                    details={"host": self.host, "port": self.port},
+                )
+            else:
+                return HealthCheckResult(
+                    status=HealthStatus.UNHEALTHY,
+                    message=f"Port {self.port} is closed",
+                    latency_ms=latency_ms,
+                    details={"host": self.host, "port": self.port, "error_code": result},
+                )
+
+        except socket.timeout:
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            return HealthCheckResult(
+                status=HealthStatus.UNHEALTHY,
+                message="Connection timed out",
+                latency_ms=latency_ms,
+                details={"host": self.host, "port": self.port},
+            )
+        except Exception as e:
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            return HealthCheckResult(
+                status=HealthStatus.UNHEALTHY,
+                message=f"Check failed: {str(e)}",
+                latency_ms=latency_ms,
+                details={"host": self.host, "port": self.port, "error": str(e)},
+            )
+
+
+class CallableHealthCheck(HealthCheck):
+    """Health check using a custom callable.
+
+    Example:
+        >>> def check_database():
+        ...     # Custom check logic
+        ...     return True, "Database OK"
+        >>>
+        >>> check = CallableHealthCheck(check_database)
+    """
+
+    def __init__(self, func: Callable[[], Union[bool, tuple]], name: str = "custom"):
+        """Initialize callable health check.
+
+        Args:
+            func: Callable that returns bool or (bool, message) tuple.
+            name: Name for this check.
+        """
+        self.func = func
+        self.name = name
+
+    def check(self) -> HealthCheckResult:
+        """Perform the custom health check."""
+        start_time = time.perf_counter()
+
+        try:
+            result = self.func()
+            latency_ms = (time.perf_counter() - start_time) * 1000
+
+            if isinstance(result, bool):
+                healthy = result
+                message = "Check passed" if healthy else "Check failed"
+            elif isinstance(result, tuple) and len(result) >= 2:
+                healthy = result[0]
+                message = result[1]
+            else:
+                healthy = bool(result)
+                message = str(result)
+
+            return HealthCheckResult(
+                status=HealthStatus.HEALTHY if healthy else HealthStatus.UNHEALTHY,
+                message=message,
+                latency_ms=latency_ms,
+                details={"check_name": self.name},
+            )
+
+        except Exception as e:
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            return HealthCheckResult(
+                status=HealthStatus.UNHEALTHY,
+                message=f"Check raised exception: {str(e)}",
+                latency_ms=latency_ms,
+                details={"check_name": self.name, "error": str(e)},
+            )
+
+
+@dataclass
+class CachedHealthResult:
+    """Cached health check result with TTL."""
+
+    result: HealthCheckResult
+    expires_at: datetime
+
+
+class HealthCheckRegistry:
+    """Registry for managing tool health checks.
+
+    Provides caching, background checking, and aggregated health status.
+
+    Example:
+        >>> registry = HealthCheckRegistry()
+        >>> registry.register("api_tool", HttpHealthCheck("https://api.example.com/health"))
+        >>>
+        >>> # Check all tools
+        >>> status = registry.check_all()
+        >>> for name, result in status.items():
+        ...     print(f"{name}: {result.status}")
+    """
+
+    def __init__(self, cache_ttl: timedelta = timedelta(seconds=30), check_timeout: float = 10.0):
+        """Initialize health check registry.
+
+        Args:
+            cache_ttl: How long to cache health check results.
+            check_timeout: Default timeout for health checks.
+        """
+        self._checks: Dict[str, HealthCheck] = {}
+        self._cache: Dict[str, CachedHealthResult] = {}
+        self._cache_ttl = cache_ttl
+        self._check_timeout = check_timeout
+        self._lock = threading.RLock()
+        self._background_thread: Optional[threading.Thread] = None
+        self._stop_event = threading.Event()
+
+    def register(self, tool_name: str, check: Union[HealthCheck, Callable[[], bool], str]) -> None:
+        """Register a health check for a tool.
+
+        Args:
+            tool_name: Name of the tool.
+            check: HealthCheck instance, callable, or URL string.
+        """
+        with self._lock:
+            if isinstance(check, str):
+                # Assume it's a URL for HTTP check
+                check = HttpHealthCheck(check)
+            elif callable(check) and not isinstance(check, HealthCheck):
+                check = CallableHealthCheck(check, name=tool_name)
+
+            self._checks[tool_name] = check
+
+    def unregister(self, tool_name: str) -> bool:
+        """Unregister a health check.
+
+        Args:
+            tool_name: Name of the tool.
+
+        Returns:
+            True if was registered, False otherwise.
+        """
+        with self._lock:
+            if tool_name in self._checks:
+                del self._checks[tool_name]
+                self._cache.pop(tool_name, None)
+                return True
+            return False
+
+    def check(self, tool_name: str, use_cache: bool = True) -> HealthCheckResult:
+        """Check health of a specific tool.
+
+        Args:
+            tool_name: Name of the tool.
+            use_cache: Whether to use cached result if available.
+
+        Returns:
+            HealthCheckResult.
+        """
+        with self._lock:
+            if tool_name not in self._checks:
+                return HealthCheckResult(
+                    status=HealthStatus.UNKNOWN,
+                    message=f"No health check registered for '{tool_name}'",
+                )
+
+            # Check cache
+            if use_cache and tool_name in self._cache:
+                cached = self._cache[tool_name]
+                if datetime.now() < cached.expires_at:
+                    return cached.result
+
+            # Perform check
+            check = self._checks[tool_name]
+
+        result = check.check()
+
+        # Update cache
+        with self._lock:
+            self._cache[tool_name] = CachedHealthResult(
+                result=result, expires_at=datetime.now() + self._cache_ttl
+            )
+
+        return result
+
+    async def check_async(self, tool_name: str, use_cache: bool = True) -> HealthCheckResult:
+        """Check health of a specific tool asynchronously.
+
+        Args:
+            tool_name: Name of the tool.
+            use_cache: Whether to use cached result if available.
+
+        Returns:
+            HealthCheckResult.
+        """
+        with self._lock:
+            if tool_name not in self._checks:
+                return HealthCheckResult(
+                    status=HealthStatus.UNKNOWN,
+                    message=f"No health check registered for '{tool_name}'",
+                )
+
+            # Check cache
+            if use_cache and tool_name in self._cache:
+                cached = self._cache[tool_name]
+                if datetime.now() < cached.expires_at:
+                    return cached.result
+
+            check = self._checks[tool_name]
+
+        result = await check.check_async()
+
+        # Update cache
+        with self._lock:
+            self._cache[tool_name] = CachedHealthResult(
+                result=result, expires_at=datetime.now() + self._cache_ttl
+            )
+
+        return result
+
+    def check_all(self, use_cache: bool = True) -> Dict[str, HealthCheckResult]:
+        """Check health of all registered tools.
+
+        Args:
+            use_cache: Whether to use cached results.
+
+        Returns:
+            Dictionary mapping tool names to results.
+        """
+        with self._lock:
+            tool_names = list(self._checks.keys())
+
+        return {name: self.check(name, use_cache) for name in tool_names}
+
+    async def check_all_async(self, use_cache: bool = True) -> Dict[str, HealthCheckResult]:
+        """Check health of all registered tools asynchronously.
+
+        Args:
+            use_cache: Whether to use cached results.
+
+        Returns:
+            Dictionary mapping tool names to results.
+        """
+        with self._lock:
+            tool_names = list(self._checks.keys())
+
+        tasks = [self.check_async(name, use_cache) for name in tool_names]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        health_results = {}
+        for name, result in zip(tool_names, results):
+            if isinstance(result, Exception):
+                health_results[name] = HealthCheckResult(
+                    status=HealthStatus.UNHEALTHY, message=f"Check failed: {str(result)}"
+                )
+            else:
+                health_results[name] = result
+
+        return health_results
+
+    def get_overall_status(self) -> HealthStatus:
+        """Get overall health status across all tools.
+
+        Returns:
+            HEALTHY if all healthy, DEGRADED if any degraded, UNHEALTHY if any unhealthy.
+        """
+        results = self.check_all()
+
+        if not results:
+            return HealthStatus.UNKNOWN
+
+        statuses = [r.status for r in results.values()]
+
+        if all(s == HealthStatus.HEALTHY for s in statuses):
+            return HealthStatus.HEALTHY
+        elif any(s == HealthStatus.UNHEALTHY for s in statuses):
+            return HealthStatus.UNHEALTHY
+        elif any(s == HealthStatus.DEGRADED for s in statuses):
+            return HealthStatus.DEGRADED
+        else:
+            return HealthStatus.UNKNOWN
+
+    def start_background_checks(self, interval: timedelta = timedelta(seconds=30)) -> None:
+        """Start background health checks.
+
+        Args:
+            interval: How often to run checks.
+        """
+        if self._background_thread is not None and self._background_thread.is_alive():
+            return
+
+        self._stop_event.clear()
+
+        def check_loop():
+            import contextlib
+
+            while not self._stop_event.is_set():
+                with contextlib.suppress(Exception):
+                    self.check_all(use_cache=False)
+                self._stop_event.wait(interval.total_seconds())
+
+        self._background_thread = threading.Thread(target=check_loop, daemon=True)
+        self._background_thread.start()
+
+    def stop_background_checks(self) -> None:
+        """Stop background health checks."""
+        self._stop_event.set()
+        if self._background_thread is not None:
+            self._background_thread.join(timeout=5.0)
+            self._background_thread = None
+
+    def clear_cache(self) -> None:
+        """Clear the health check cache."""
+        with self._lock:
+            self._cache.clear()
+
+
+# Global health check registry
+_global_health_registry: HealthCheckRegistry = HealthCheckRegistry()
+
+
+def get_health_registry() -> HealthCheckRegistry:
+    """Get the global health check registry.
+
+    Returns:
+        The global HealthCheckRegistry instance.
+    """
+    return _global_health_registry
+
+
+def set_health_registry(registry: HealthCheckRegistry) -> None:
+    """Set the global health check registry.
+
+    Args:
+        registry: The registry to use globally.
+    """
+    global _global_health_registry
+    _global_health_registry = registry