from __future__ import annotations import asyncio import ssl import time from collections.abc import Callable from yuxi.channels.models import HealthStatus async def ping_probe( send_fn: Callable[[str], None], read_fn, server: str, timeout: float = 8.0, metadata: dict | None = None, ) -> HealthStatus: start = time.monotonic() try: ping_token = f"fp-health-{int(time.time())}" send_fn(f"PING :{ping_token}") deadline = time.monotonic() + timeout while time.monotonic() < deadline: try: line = await asyncio.wait_for(read_fn(), timeout=deadline - time.monotonic()) except TimeoutError: break if line and "PONG" in line: latency_ms = (time.monotonic() - start) * 1000 return HealthStatus( status="healthy", latency_ms=latency_ms, metadata=metadata or {}, ) return HealthStatus( status="degraded", last_error="PING no PONG response within timeout", ) except (ConnectionError, OSError) as e: return HealthStatus(status="unhealthy", last_error=str(e)) except Exception as e: return HealthStatus(status="unhealthy", last_error=str(e)) async def standalone_tcp_probe( host: str, port: int = 6697, use_tls: bool = True, nick: str = "ForcePilotProbe", username: str = "probe", realname: str = "ForcePilot Probe", timeout: float = 15.0, ) -> HealthStatus: """Independent TCP probe that connects, waits for 001, and disconnects.""" start = time.monotonic() reader: asyncio.StreamReader | None = None writer: asyncio.StreamWriter | None = None try: if use_tls: ctx = ssl.create_default_context() ctx.check_hostname = True ctx.verify_mode = ssl.CERT_REQUIRED reader, writer = await asyncio.wait_for( asyncio.open_connection( host=host, port=port, ssl=ctx, server_hostname=host, ), timeout=timeout, ) else: reader, writer = await asyncio.wait_for( asyncio.open_connection(host=host, port=port), timeout=timeout, ) caps = ["server-time"] writer.write(b"CAP LS 302\r\n") writer.write(f"CAP REQ :{' '.join(caps)}\r\n".encode()) writer.write(f"NICK {nick}\r\n".encode()) writer.write(f"USER {username} 0 * :{realname}\r\n".encode()) await writer.drain() registered = False deadline = time.monotonic() + timeout while time.monotonic() < deadline: line = await asyncio.wait_for(reader.readline(), timeout=deadline - time.monotonic()) if not line: break decoded = line.decode("utf-8", errors="replace").rstrip("\r\n") parts = decoded.split() if len(parts) >= 2 and parts[1] in ("001", "376", "422"): registered = True break if "ERROR" in decoded: return HealthStatus( status="unhealthy", latency_ms=(time.monotonic() - start) * 1000, metadata={"host": host, "port": port, "tls": use_tls}, last_error=f"Server returned: {decoded}", ) if not registered: return HealthStatus( status="degraded", latency_ms=(time.monotonic() - start) * 1000, metadata={"host": host, "port": port, "tls": use_tls}, last_error="Registration not completed within timeout", ) latency_ms = (time.monotonic() - start) * 1000 try: writer.write(b"QUIT :ForcePilot probe done\r\n") await writer.drain() except Exception: pass return HealthStatus( status="healthy", latency_ms=latency_ms, metadata={"host": host, "port": port, "tls": use_tls, "probe_type": "standalone_tcp"}, ) except TimeoutError: return HealthStatus( status="unhealthy", latency_ms=(time.monotonic() - start) * 1000, metadata={"host": host, "port": port, "tls": use_tls}, last_error="Connection timed out", ) except (ConnectionError, OSError) as e: return HealthStatus( status="unhealthy", latency_ms=(time.monotonic() - start) * 1000, metadata={"host": host, "port": port, "tls": use_tls}, last_error=str(e), ) except Exception as e: return HealthStatus( status="unhealthy", latency_ms=(time.monotonic() - start) * 1000, metadata={"host": host, "port": port, "tls": use_tls}, last_error=str(e), ) finally: if writer: try: writer.close() except Exception: pass