from __future__ import annotations import asyncio import json import random from collections.abc import Awaitable, Callable import websockets from yuxi.channels.infra.circuit_breaker import CircuitBreaker, CircuitBreakerOpenError from yuxi.utils.logging_config import logger class RelayManager: def __init__(self, relay_urls: list[str], timeout: float = 30.0, reconnect_interval: float = 5.0): self._relay_urls = relay_urls self._timeout = timeout self._reconnect_interval = reconnect_interval self._connections: dict[str, websockets.WebSocketClientProtocol] = {} self._event_handlers: list[Callable[[dict], Awaitable[None]]] = [] self._eose_handlers: list[Callable[[str], Awaitable[None]]] = [] self._reconnect_tasks: dict[str, asyncio.Task] = {} self._subscriptions: list[dict] = [] self._running = False self._circuit_breakers: dict[str, CircuitBreaker] = {} self._on_connect_handlers: list[Callable[[str], Awaitable[None]]] = [] self._on_disconnect_handlers: list[Callable[[str], Awaitable[None]]] = [] self._relay_scores: dict[str, float] = {} self._on_error_handlers: list[Callable[[str, dict], Awaitable[None]]] = [] self._reconnect_attempts: dict[str, int] = {} for url in relay_urls: self._circuit_breakers[url] = CircuitBreaker( failure_threshold=5, recovery_timeout=30.0, half_open_max_calls=3 ) self._relay_scores[url] = 0.5 self._sub_pubkeys: list[str] = [] def set_pubkey_filter(self, pubkeys: list[str]) -> None: self._sub_pubkeys = pubkeys def _build_subscription_filters(self, filters: list[dict]) -> list[dict]: if not self._sub_pubkeys: return filters return [{**f, "#p": self._sub_pubkeys} for f in filters] def on_connect(self, handler: Callable[[str], Awaitable[None]]) -> None: self._on_connect_handlers.append(handler) def on_disconnect(self, handler: Callable[[str], Awaitable[None]]) -> None: self._on_disconnect_handlers.append(handler) def on_eose(self, handler: Callable[[str], Awaitable[None]]) -> None: self._eose_handlers.append(handler) def on_error(self, handler: Callable[[str, dict], Awaitable[None]]) -> None: self._on_error_handlers.append(handler) def _classify_error(self, url: str, error: Exception) -> str: error_type = type(error).__name__ if "timeout" in error_type.lower() or "Timeout" in error_type: return "timeout" if "connect" in error_type.lower() or "Connection" in error_type: return "connection" if "circuit" in error_type.lower(): return "circuit_breaker" if "json" in error_type.lower(): return "parse_error" if "websocket" in error_type.lower(): return "websocket" return "unknown" async def _notify_error(self, url: str, category: str, error: Exception) -> None: context = { "url": url, "category": category, "error_type": type(error).__name__, "error_message": str(error)[:200], } for handler in self._on_error_handlers: try: await handler(category, context) except Exception: logger.debug(f"on_error handler error for {url}", exc_info=True) async def connect_all(self) -> None: self._running = True tasks = [self._connect_relay(url) for url in self._relay_urls] results = await asyncio.gather(*tasks, return_exceptions=True) connected = sum(1 for r in results if r is True) logger.info(f"Nostr Relay 连接完成: {connected}/{len(self._relay_urls)} 个 Relay 已连接") async def _connect_relay(self, url: str) -> bool: try: ws = await asyncio.wait_for( websockets.connect(url, ping_interval=20, ping_timeout=10), timeout=self._timeout, ) self._connections[url] = ws if self._subscriptions: sub_id = "forcepilot_nostr_sub" effective_filters = self._build_subscription_filters(self._subscriptions) req = json.dumps(["REQ", sub_id, *effective_filters]) await ws.send(req) for handler in self._on_connect_handlers: try: await handler(url) except Exception: logger.debug(f"on_connect handler error for {url}", exc_info=True) logger.info(f"Nostr Relay 已连接: {url}") return True except Exception as e: logger.warning(f"Nostr Relay 连接失败 {url}: {e}") self._schedule_reconnect(url) return False def _schedule_reconnect(self, url: str) -> None: if url in self._reconnect_tasks: return attempt = self._reconnect_attempts.get(url, 0) + 1 self._reconnect_attempts[url] = attempt base_delay = min(self._reconnect_interval * (2 ** (attempt - 1)), 300) jitter = random.uniform(0, base_delay * 0.3) delay = base_delay + jitter async def _reconnect(): await asyncio.sleep(delay) if not self._running: self._reconnect_tasks.pop(url, None) return logger.debug(f"Nostr Relay 尝试重连: {url} (尝试 #{attempt}, 延迟 {delay:.1f}s)") success = await self._connect_relay(url) if success: self._reconnect_attempts.pop(url, None) self._reconnect_tasks.pop(url, None) self._reconnect_tasks[url] = asyncio.ensure_future(_reconnect()) def list_connections(self) -> dict[str, websockets.WebSocketClientProtocol]: return dict(self._connections) def get_circuit_breaker(self, url: str) -> CircuitBreaker | None: return self._circuit_breakers.get(url) async def disconnect_all(self) -> None: self._running = False for url, task in list(self._reconnect_tasks.items()): task.cancel() self._reconnect_tasks.pop(url, None) for url in list(self._connections.keys()): ws = self._connections.pop(url, None) if ws: try: await ws.close() except Exception: pass for handler in self._on_disconnect_handlers: try: await handler(url) except Exception: logger.debug(f"on_disconnect handler error for {url}", exc_info=True) async def broadcast(self, event: dict) -> int: message = json.dumps(["EVENT", event]) sorted_urls = sorted( self._connections.keys(), key=lambda url: self._relay_scores.get(url, 0.5), reverse=True, ) tasks = [] for url in sorted_urls: cb = self._circuit_breakers.get(url) if cb and cb.state == "open": continue ws = self._connections.get(url) if ws and ws.open: tasks.append(self._send_with_cb(ws, url, message)) if not tasks: return 0 results = await asyncio.gather(*tasks, return_exceptions=True) return sum(1 for r in results if r is True) async def broadcast_with_ack(self, event: dict) -> dict: message = json.dumps(["EVENT", event]) sorted_urls = sorted( self._connections.keys(), key=lambda url: self._relay_scores.get(url, 0.5), reverse=True, ) ack_results: dict[str, bool] = {} tasks = [] for url in sorted_urls: cb = self._circuit_breakers.get(url) if cb and cb.state == "open": ack_results[url] = False continue ws = self._connections.get(url) if ws and ws.open: tasks.append(self._send_and_wait_ok(ws, url, message, ack_results)) if tasks: await asyncio.gather(*tasks, return_exceptions=True) return ack_results async def _send_and_wait_ok(self, ws, url: str, message: str, results: dict[str, bool]) -> None: try: await ws.send(message) raw = await asyncio.wait_for(ws.recv(), timeout=5) data = json.loads(raw) ok = isinstance(data, list) and len(data) >= 4 and data[0] == "OK" and data[2] is True results[url] = ok except Exception: results[url] = False async def _send_with_cb( self, ws: websockets.WebSocketClientProtocol, url: str, message: str, max_retries: int = 3 ) -> bool: cb = self._circuit_breakers.get(url) retry_delays = [1, 2, 4] for attempt in range(max_retries): try: if cb: await cb.call(lambda: asyncio.wait_for(ws.send(message), timeout=10)) else: await asyncio.wait_for(ws.send(message), timeout=10) return True except CircuitBreakerOpenError: logger.debug(f"Relay {url} circuit breaker open, skipping") return False except Exception as e: if cb: await cb.record_failure() category = self._classify_error(url, e) await self._notify_error(url, category, e) if attempt < max_retries - 1: await asyncio.sleep(retry_delays[min(attempt, len(retry_delays) - 1)]) logger.warning(f"Relay 发送失败(已重试{max_retries}次): {url}，断开并调度重连") self._connections.pop(url, None) self._schedule_reconnect(url) return False async def _send_with_retry( self, ws: websockets.WebSocketClientProtocol, url: str, message: str, max_retries: int = 3 ) -> bool: return await self._send_with_cb(ws, url, message, max_retries) def on_event(self, handler: Callable[[dict], Awaitable[None]]) -> None: self._event_handlers.append(handler) async def subscribe(self, filters: list[dict]) -> None: self._subscriptions = filters sub_id = "forcepilot_nostr_sub" effective_filters = self._build_subscription_filters(filters) req = json.dumps(["REQ", sub_id, *effective_filters]) for url in list(self._connections.keys()): ws = self._connections.get(url) if ws and ws.open: try: await ws.send(req) except Exception: pass async def listen(self, url: str) -> None: ws = self._connections.get(url) if not ws or not ws.open: return try: async for raw in ws: if not self._running: break try: data = json.loads(raw) if isinstance(data, list) and len(data) >= 3 and data[0] == "EVENT": event = data[2] if isinstance(event, dict): for handler in self._event_handlers: await handler(event) elif isinstance(data, list) and len(data) >= 2 and data[0] == "EOSE": for handler in self._eose_handlers: try: await handler(url) except Exception: logger.debug(f"EOSE handler error for {url}", exc_info=True) except (json.JSONDecodeError, KeyError, IndexError): logger.debug(f"Relay {url} 返回无法解析的消息: {raw[:200]}") continue except Exception: self._connections.pop(url, None) for handler in self._on_disconnect_handlers: try: await handler(url) except Exception: pass self._schedule_reconnect(url) async def query(self, filters: list[dict], timeout: float = 10.0) -> list[dict]: sub_id = "forcepilot_nostr_query" req = json.dumps(["REQ", sub_id, *filters]) close_req = json.dumps(["CLOSE", sub_id]) active_ws = [(url, ws) for url, ws in self._connections.items() if ws and ws.open] if not active_ws: return [] for url, ws in active_ws: try: await ws.send(req) except Exception: continue async def _receive_from_relay(_url: str, _ws) -> list[dict]: results: list[dict] = [] try: while True: raw = await asyncio.wait_for(_ws.recv(), timeout=timeout) try: data = json.loads(raw) except json.JSONDecodeError: continue if isinstance(data, list) and len(data) >= 2: if data[0] == "EVENT" and len(data) >= 3 and isinstance(data[2], dict): results.append(data[2]) elif data[0] in ("EOSE", "CLOSED"): break except TimeoutError: pass except Exception: pass return results receive_tasks = [_receive_from_relay(url, ws) for url, ws in active_ws] all_results = await asyncio.gather(*receive_tasks) seen: set[str] = set() events: list[dict] = [] for batch in all_results: for event in batch: event_id = event.get("id", "") if event_id and event_id not in seen: seen.add(event_id) events.append(event) for url, ws in active_ws: try: await ws.send(close_req) except Exception: pass return events async def send_auth(self, raw_signed_auth_event: dict, urls: list[str] | None = None) -> dict[str, bool]: message = json.dumps(["AUTH", raw_signed_auth_event]) target_urls = urls or list(self._connections.keys()) results: dict[str, bool] = {} for url in target_urls: ws = self._connections.get(url) if not ws or not ws.open: results[url] = False continue try: await ws.send(message) raw = await asyncio.wait_for(ws.recv(), timeout=10) data = json.loads(raw) results[url] = (isinstance(data, list) and len(data) >= 2 and data[0] == "AUTH" and data[1] == url) or ( isinstance(data, list) and len(data) >= 2 and data[0] == "OK" ) except Exception: results[url] = False return results def active_count(self) -> tuple[int, int]: active = sum(1 for ws in self._connections.values() if ws.open) return active, len(self._relay_urls) def set_relay_scores(self, scores: dict[str, float]) -> None: self._relay_scores.update(scores) def get_relay_score(self, url: str) -> float: return self._relay_scores.get(url, 0.5)