ForcePilot/backend/package/yuxi/channel/gateway/probe.py

import asyncio
import logging
import time
from dataclasses import dataclass, field
from enum import StrEnum

import aiohttp

from yuxi.channel.gateway.protocol import (
    RpcRequest,
    RpcResponse,
    marshal_frame,
    unmarshal_frame,
)

logger = logging.getLogger(__name__)

DEFAULT_PROBE_TIMEOUT_MS = 10_000
MIN_PROBE_TIMEOUT_MS = 250
MAX_PROBE_TIMEOUT_MS = 30_000


class ProbeCapability(StrEnum):
    UNKNOWN = "unknown"
    PAIRING_PENDING = "pairing_pending"
    CONNECTED_NO_OPERATOR_SCOPE = "connected_no_operator_scope"
    READ_ONLY = "read_only"
    WRITE_CAPABLE = "write_capable"
    ADMIN_CAPABLE = "admin_capable"


@dataclass
class ProbeAuthSummary:
    role: str | None = None
    scopes: list[str] = field(default_factory=list)
    capability: ProbeCapability = ProbeCapability.UNKNOWN


@dataclass
class ProbeServerSummary:
    version: str | None = None
    conn_id: str | None = None


@dataclass
class ProbeClose:
    code: int
    reason: str
    hint: str | None = None


@dataclass
class GatewayProbeResult:
    ok: bool
    url: str
    connect_latency_ms: float | None = None
    error: str | None = None
    close: ProbeClose | None = None
    auth: ProbeAuthSummary = field(default_factory=ProbeAuthSummary)
    server: ProbeServerSummary = field(default_factory=ProbeServerSummary)
    health: dict | None = None
    status: dict | None = None
    presence: list | None = None
    config_snapshot: dict | None = None


def _clamp_timeout_ms(timeout_ms: float | int) -> float:
    return max(MIN_PROBE_TIMEOUT_MS, min(float(timeout_ms), MAX_PROBE_TIMEOUT_MS))


def _resolve_capability(
    scopes: list[str],
    connect_latency_ms: float | None,
    auth_metadata_present: bool,
) -> ProbeCapability:
    if "operator.admin" in scopes:
        return ProbeCapability.ADMIN_CAPABLE
    if "operator.write" in scopes:
        return ProbeCapability.WRITE_CAPABLE
    if "operator.read" in scopes or "admin" in scopes:
        return ProbeCapability.READ_ONLY
    if connect_latency_ms is not None and auth_metadata_present:
        return ProbeCapability.CONNECTED_NO_OPERATOR_SCOPE
    return ProbeCapability.UNKNOWN


async def probe_gateway(
    url: str,
    token: str | None = None,
    password: str | None = None,
    timeout_ms: float = DEFAULT_PROBE_TIMEOUT_MS,
    detail_level: str = "full",
) -> GatewayProbeResult:
    connect_latency_ms: float | None = None
    connect_error: str | None = None
    close_info: ProbeClose | None = None
    auth_summary = ProbeAuthSummary()
    server_summary = ProbeServerSummary()
    auth_metadata_present = False

    effective_timeout = _clamp_timeout_ms(timeout_ms)
    timeout_sec = effective_timeout / 1000.0

    headers: dict[str, str] = {}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    elif password:
        headers["Authorization"] = f"Bearer {password}"

    conn_start = time.monotonic()

    try:
        async with aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=timeout_sec),
        ) as session:
            async with session.ws_connect(
                url,
                headers=headers,
                heartbeat=15.0,
            ) as ws:
                connect_latency_ms = (time.monotonic() - conn_start) * 1000

                auth_metadata_present = True

                if detail_level == "none":
                    ws_close = ws.close_code
                    if ws_close is not None:
                        close_info = ProbeClose(
                            code=ws_close,
                            reason="connection closed by server",
                        )
                    return GatewayProbeResult(
                        ok=ws_close is None,
                        url=url,
                        connect_latency_ms=connect_latency_ms,
                        error=None,
                        close=close_info,
                        auth=auth_summary,
                        server=server_summary,
                    )

                rpc_request_id = "probe"
                rpc_request = RpcRequest(
                    id=rpc_request_id,
                    method="system.health",
                    params={},
                )
                await ws.send_str(marshal_frame(rpc_request))

                try:
                    raw = await asyncio.wait_for(
                        ws.receive_str(),
                        timeout=effective_timeout / 1000.0,
                    )
                except asyncio.TimeoutError:
                    return GatewayProbeResult(
                        ok=False,
                        url=url,
                        connect_latency_ms=connect_latency_ms,
                        error="timeout waiting for health response",
                        auth=auth_summary,
                        server=server_summary,
                    )

                try:
                    frame = unmarshal_frame(raw)
                except Exception:
                    return GatewayProbeResult(
                        ok=False,
                        url=url,
                        connect_latency_ms=connect_latency_ms,
                        error="invalid response frame from gateway",
                        auth=auth_summary,
                        server=server_summary,
                    )

                if not isinstance(frame, RpcResponse) or not frame.ok:
                    error_msg = frame.error_message if isinstance(frame, RpcResponse) else "unexpected frame type"
                    return GatewayProbeResult(
                        ok=False,
                        url=url,
                        connect_latency_ms=connect_latency_ms,
                        error=error_msg or "gateway returned error",
                        auth=auth_summary,
                        server=server_summary,
                    )

                health_data = frame.result

                auth_summary = ProbeAuthSummary(
                    role="viewer",
                    scopes=[],
                    capability=ProbeCapability.READ_ONLY,
                )

                if detail_level == "presence":
                    return GatewayProbeResult(
                        ok=True,
                        url=url,
                        connect_latency_ms=connect_latency_ms,
                        health=health_data,
                        auth=auth_summary,
                        server=server_summary,
                        presence=health_data.get("channels") if health_data else None,
                    )

                status_data: dict | None = None
                try:
                    status_request = RpcRequest(
                        id="probe-status",
                        method="channels.status",
                        params={},
                    )
                    await ws.send_str(marshal_frame(status_request))
                    raw_status = await asyncio.wait_for(
                        ws.receive_str(),
                        timeout=effective_timeout / 1000.0,
                    )
                    status_frame = unmarshal_frame(raw_status)
                    if isinstance(status_frame, RpcResponse) and status_frame.ok:
                        status_data = status_frame.result
                except Exception:
                    logger.debug("Failed to get channel status during probe", exc_info=True)

                return GatewayProbeResult(
                    ok=True,
                    url=url,
                    connect_latency_ms=connect_latency_ms,
                    health=health_data,
                    status=status_data,
                    auth=auth_summary,
                    server=server_summary,
                )

    except aiohttp.ClientConnectorError as e:
        connect_error = f"连接失败: {e}"
    except aiohttp.WSServerHandshakeError as e:
        connect_error = f"WebSocket 握手失败: {e.status} {e.message}"
    except aiohttp.ClientError as e:
        connect_error = f"客户端错误: {e}"
    except asyncio.TimeoutError:
        connect_error = "连接超时"
    except Exception as e:
        connect_error = f"未知错误: {e}"

    return GatewayProbeResult(
        ok=False,
        url=url,
        connect_latency_ms=connect_latency_ms,
        error=connect_error,
        close=close_info,
        auth=auth_summary,
        server=server_summary,
    )


def format_probe_result(result: GatewayProbeResult, verbose: bool = False) -> str:
    lines: list[str] = []

    if result.ok:
        lines.append(f"✓ 网关探测成功: {result.url}")
        if result.connect_latency_ms is not None:
            lines.append(f"  连接延迟: {result.connect_latency_ms:.1f}ms")
        lines.append(f"  能力级别: {result.auth.capability.value}")
    else:
        lines.append(f"✗ 网关探测失败: {result.url}")
        if result.connect_latency_ms is not None:
            lines.append(f"  连接延迟: {result.connect_latency_ms:.1f}ms")
        if result.error:
            lines.append(f"  错误: {result.error}")
        if result.close:
            lines.append(f"  关闭信息: code={result.close.code} reason={result.close.reason}")

    if verbose and result.health:
        lines.append("")
        lines.append("健康报告:")
        health = result.health
        if isinstance(health, dict):
            lines.append(f"  状态: {health.get('status', 'unknown')}")
            summary = health.get("summary", {})
            if summary:
                lines.append(
                    f"  频道: 总计={summary.get('total', 0)}, "
                    f"运行={summary.get('running', 0)}, "
                    f"停止={summary.get('stopped', 0)}, "
                    f"异常={summary.get('unhealthy', 0)}"
                )

    return "\n".join(lines)