ForcePilot/backend/package/yuxi/channels/adapters/qqbot/probe.py

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)