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)