import asyncio import json import logging import time import uuid from collections.abc import Callable import httpx from yuxi.channel.extensions.tlon.errors import UrbitError, UrbitSSEError from yuxi.channel.extensions.tlon.types import SSEConfig logger = logging.getLogger(__name__) class UrbitSSEClient: def __init__(self, config: SSEConfig): self._config = config self._channel_id = f"{int(time.time() * 1000)}-{uuid.uuid4().hex[:8]}" self._subscriptions: dict[int, dict] = {} self._event_handlers: dict[int, callable] = {} self._event_count = 0 self._state = "idle" self._abort = asyncio.Event() self._http: httpx.AsyncClient | None = None @property def channel_id(self) -> str: return self._channel_id @property def state(self) -> str: return self._state async def connect(self) -> None: self._http = self._create_http_client() await self._create_channel() await self._poke_helm_hi() await self._subscribe_all() self._state = "connected" async def subscribe(self, app: str, path: str, on_event: Callable, on_error: Callable | None = None, on_quit: Callable | None = None) -> int: sub_id = len(self._subscriptions) + 1 self._subscriptions[sub_id] = { "app": app, "path": path, "on_error": on_error, "on_quit": on_quit, } self._event_handlers[sub_id] = on_event if self._state == "connected": await self._poke_subscribe(sub_id, app, path) return sub_id async def poke(self, app: str, mark: str, json_data: dict) -> int: poke_id = int(time.time() * 1000) payload = [{ "id": poke_id, "action": "poke", "ship": self._config.ship, "app": app, "mark": mark, "json": json_data, }] await self._put_channel(payload) return poke_id async def scry(self, path: str) -> dict: scry_url = f"{self._config.url}/~/scry{path}" async with httpx.AsyncClient(timeout=15.0) as client: response = await client.get( scry_url, headers={"Cookie": self._config.cookie}, ) response.raise_for_status() return response.json() async def listen(self) -> None: if self._state != "connected": raise UrbitSSEError("Client not connected") async with self._http.stream( "GET", f"/~/channel/{self._channel_id}", headers={ "Cookie": self._config.cookie, "Accept": "text/event-stream", }, timeout=self._config.connect_timeout, ) as response: self._state = "streaming" event_id = None data_lines = [] async for line in response.aiter_lines(): if self._abort.is_set(): break if line.startswith("id:"): event_id = int(line[3:].strip()) elif line.startswith("data:"): data_lines.append(line[5:]) elif line == "" and event_id is not None: await self._handle_event(event_id, data_lines) event_id = None data_lines = [] elif line == "": event_id = None data_lines = [] async def close(self) -> None: self._state = "closing" try: for sub_id in self._subscriptions: await self._poke_unsubscribe(sub_id) if self._http: await self._http.delete( f"/~/channel/{self._channel_id}", headers={"Cookie": self._config.cookie}, ) finally: self._abort.set() if self._http: await self._http.aclose() self._state = "closed" async def attempt_reconnect(self) -> None: logger.info("[tlon] SSE reconnecting...") for attempt in range(1, self._config.max_reconnect_attempts + 1): if self._abort.is_set(): return delay = min( self._config.reconnect_delay_ms * (2 ** (attempt - 1)), self._config.max_reconnect_delay_ms, ) await asyncio.sleep(delay / 1000) try: self._config.cookie = await self._config.on_reconnect() except Exception: continue self._channel_id = f"{int(time.time() * 1000)}-{uuid.uuid4().hex[:8]}" try: if self._http: await self._http.aclose() await self.connect() return except Exception: continue logger.warning("[tlon] Max reconnect attempts reached, retrying after delay") await asyncio.sleep(10) await self.attempt_reconnect() async def abort(self) -> None: self._abort.set() async def _handle_event(self, event_id: int, data_lines: list[str]) -> None: self._event_count += 1 raw_data = "".join(data_lines) try: data = json.loads(raw_data) except json.JSONDecodeError: return if isinstance(data, dict) and data.get("response") == "quit": for sub_id, info in self._subscriptions.items(): if info.get("on_quit"): await info["on_quit"](data) return for sub_id, handler in self._event_handlers.items(): try: await handler(data) except Exception: logger.exception("[tlon] Event handler error") if self._event_count % self._config.ack_threshold == 0: await self._send_ack(event_id) async def _send_ack(self, event_id: int) -> None: try: await self._http.put( f"/~/channel/{self._channel_id}", json=[{"id": event_id, "action": "ack"}], headers={"Cookie": self._config.cookie}, ) except Exception: pass def _create_http_client(self) -> httpx.AsyncClient: return httpx.AsyncClient( base_url=self._config.url, timeout=30.0, ) async def _create_channel(self) -> None: response = await self._http.put( f"/~/channel/{self._channel_id}", json=[], headers={"Cookie": self._config.cookie, "Content-Type": "application/json"}, ) if response.status_code not in (200, 204): raise UrbitError(f"Channel PUT failed: HTTP {response.status_code}") async def _poke_helm_hi(self) -> None: await self._http.put( f"/~/channel/{self._channel_id}", json=[{ "id": 1, "action": "poke", "ship": self._config.ship, "app": "hood", "mark": "helm-hi", "json": "opening airlock", }], headers={"Cookie": self._config.cookie, "Content-Type": "application/json"}, ) async def _poke_subscribe(self, sub_id: int, app: str, path: str) -> None: await self._http.put( f"/~/channel/{self._channel_id}", json=[{ "id": sub_id, "action": "subscribe", "ship": self._config.ship, "app": app, "path": path, }], headers={"Cookie": self._config.cookie, "Content-Type": "application/json"}, ) async def _poke_unsubscribe(self, sub_id: int) -> None: try: sub = self._subscriptions.get(sub_id) if sub: await self._http.put( f"/~/channel/{self._channel_id}", json=[{ "id": sub_id, "action": "unsubscribe", "ship": self._config.ship, "app": sub["app"], "path": sub["path"], }], headers={"Cookie": self._config.cookie, "Content-Type": "application/json"}, ) except Exception: pass async def _subscribe_all(self) -> None: for sub_id, sub in self._subscriptions.items(): await self._poke_subscribe(sub_id, sub["app"], sub["path"]) async def _put_channel(self, payload: list[dict]) -> None: response = await self._http.put( f"/~/channel/{self._channel_id}", json=payload, headers={"Cookie": self._config.cookie, "Content-Type": "application/json"}, ) if response.status_code not in (200, 204): raise UrbitError(f"Channel PUT failed: HTTP {response.status_code}") async def scry_static(self, path: str) -> dict: return await self.scry(path)