267 lines
9.0 KiB
Python
267 lines
9.0 KiB
Python
|
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)
|