from __future__ import annotations import time from dataclasses import dataclass, field import aiohttp from yuxi.channels.models import HealthStatus from yuxi.utils.logging_config import logger @dataclass class EndpointCheck: name: str url: str status: str = "unknown" latency_ms: float = 0 error: str = "" http_status: int = 0 @dataclass class AggregatedHealth: status: str = "healthy" overall_latency_ms: float = 0 checks: list[EndpointCheck] = field(default_factory=list) ws_latency_ms: float | None = None @property def all_healthy(self) -> bool: return all(c.status == "healthy" for c in self.checks) async def health_check_dsm( api_base: str, token: str, http_client: aiohttp.ClientSession | None = None, sandbox: bool = False, ws_connected: bool = False, ) -> HealthStatus: headers = {"Authorization": f"QQBot {token}"} timeout = aiohttp.ClientTimeout(total=10) async def _check(session: aiohttp.ClientSession) -> HealthStatus: start = time.monotonic() async with session.get(f"{api_base}/gateway", headers=headers) as resp: latency_ms = (time.monotonic() - start) * 1000 if resp.status == 200: return HealthStatus( status="healthy", latency_ms=latency_ms, metadata={ "sandbox": sandbox, "ws_connected": ws_connected, }, ) elif resp.status == 401: return HealthStatus( status="unhealthy", last_error="Token expired or invalid", metadata={ "sandbox": sandbox, "auth_status": "failed", }, ) else: return HealthStatus( status="degraded", latency_ms=latency_ms, last_error=f"Gateway returned {resp.status}", ) if http_client: try: return await _check(http_client) except Exception as e: logger.warning(f"[QQBot] Health check failed: {e}") return HealthStatus(status="unhealthy", last_error=str(e), metadata={"sandbox": sandbox}) try: async with aiohttp.ClientSession(timeout=timeout) as session: return await _check(session) except Exception as e: logger.warning(f"[QQBot] Health check failed: {e}") return HealthStatus(status="unhealthy", last_error=str(e), metadata={"sandbox": sandbox}) async def health_check_multi_endpoint( api_base: str, token: str, http_client: aiohttp.ClientSession | None = None, ws_connected: bool = False, ws_latency_ms: float | None = None, ) -> AggregatedHealth: headers = {"Authorization": f"QQBot {token}"} timeout = aiohttp.ClientTimeout(total=10) endpoints = [ EndpointCheck(name="gateway", url=f"{api_base}/gateway"), EndpointCheck(name="bot_info", url=f"{api_base}/v2/users/@me"), ] async def _check(session: aiohttp.ClientSession) -> AggregatedHealth: for ep in endpoints: try: start = time.monotonic() async with session.get(ep.url, headers=headers) as resp: ep.latency_ms = (time.monotonic() - start) * 1000 ep.http_status = resp.status if resp.status == 200: ep.status = "healthy" elif resp.status == 401: ep.status = "unhealthy" ep.error = "Authentication failed" elif resp.status >= 500: ep.status = "degraded" ep.error = f"Server error ({resp.status})" else: ep.status = "degraded" ep.error = f"Unexpected status ({resp.status})" except TimeoutError: ep.status = "degraded" ep.error = "Timeout" except Exception as e: ep.status = "unhealthy" ep.error = str(e) overall = "healthy" if any(ep.status == "unhealthy" for ep in endpoints): overall = "unhealthy" elif any(ep.status == "degraded" for ep in endpoints): overall = "degraded" healthy_checks = [ep for ep in endpoints if ep.status == "healthy"] avg_latency = sum(ep.latency_ms for ep in healthy_checks) / len(healthy_checks) if healthy_checks else 0 return AggregatedHealth( status=overall, overall_latency_ms=avg_latency, checks=endpoints, ws_latency_ms=ws_latency_ms, ) if http_client: return await _check(http_client) async with aiohttp.ClientSession(timeout=timeout) as session: return await _check(session)