ForcePilot/backend/package/yuxi/channel/extensions/tlon/gateway.py

import asyncio
import logging

from yuxi.channel.extensions.tlon.auth import authenticate
from yuxi.channel.extensions.tlon.sse_client import UrbitSSEClient, SSEConfig
from yuxi.channel.extensions.tlon.config import TlonConfigAdapter
from yuxi.channel.extensions.tlon.errors import UrbitAuthError
from yuxi.channel.extensions.tlon.monitor import monitor_tlon_provider
from yuxi.channel.extensions.tlon.doctor import TlonDoctor

logger = logging.getLogger(__name__)

_config_adapter = TlonConfigAdapter()


async def start_tlon_gateway(ctx) -> asyncio.Task | None:
    account_id = ctx.account_id
    config = ctx.config
    account = _config_adapter.resolve_account(account_id, config)

    if not _config_adapter.is_configured(account):
        logger.warning("[tlon] Account %s not configured, skipping", account_id)
        return None

    doctor = TlonDoctor()
    if doctor.check_migration_needed(config):
        ctx.config = doctor.normalize_compatibility_config(config)
        account = _config_adapter.resolve_account(account_id, ctx.config)

    ship = account["ship"]
    url = account["url"]
    code = account["code"]

    cookie = await _authenticate_with_retry(url, code)

    async def on_reconnect():
        return await authenticate(url, code)

    sse_config = SSEConfig(
        url=url,
        cookie=cookie,
        ship=ship,
        on_reconnect=on_reconnect,
    )
    client = UrbitSSEClient(sse_config)

    try:
        await client.connect()
        logger.info("[tlon] SSE client connected for %s (%s)", ship, account_id)
    except Exception as e:
        logger.error("[tlon] Failed to connect SSE client: %s", e)
        return None

    try:
        blocked = await client.scry("/chat/blocked.json")
        client._blocked_ships = blocked.get("blocked", [])
        logger.info("[tlon] Loaded %d blocked ships", len(client._blocked_ships))
    except Exception:
        client._blocked_ships = []

    try:
        from yuxi.channel.extensions.tlon.settings_store import (
            TlonSettingsStore, migrate_config_to_settings,
        )
        settings = TlonSettingsStore()
        store_data = await settings.load_from_scry(client)

        if not store_data:
            migrate_config_to_settings(config, client)
            logger.info("[tlon] Config migrated to Settings Store")
            store_data = await settings.load_from_scry(client)

        if store_data:
            from yuxi.channel.extensions.tlon.monitor import apply_settings_to_account
            await apply_settings_to_account(account, settings)
    except Exception as e:
        logger.warning("[tlon] Settings store init skipped: %s", e)

    try:
        from yuxi.channel.extensions.tlon.discovery import discover_and_track
        auto_discover = account.get("auto_discover_channels", False)
        if auto_discover or not account.get("group_channels"):
            watched = set(account.get("group_channels", []))
            discovered = await discover_and_track(client, watched)
            if discovered:
                logger.info("[tlon] Startup discovery found %d channels", len(discovered))
                account["group_channels"] = list(watched)
    except Exception as e:
        logger.warning("[tlon] Channel discovery skipped: %s", e)

    ctx._tlon_client = client
    ctx._tlon_account = account

    async def run():
        try:
            while not ctx.cancel_event.is_set():
                try:
                    await monitor_tlon_provider(ctx, client, account)
                except Exception as e:
                    logger.error("[tlon] Monitor crashed: %s", e)
                    await client.attempt_reconnect()
        finally:
            await client.close()

    task = asyncio.create_task(run())
    return task


async def stop_tlon_gateway(ctx) -> None:
    client = getattr(ctx, "_tlon_client", None)
    if client:
        await client.close()


async def _authenticate_with_retry(url: str, code: str, max_attempts: int = 10) -> str:
    for attempt in range(1, max_attempts + 1):
        try:
            return await authenticate(url, code)
        except UrbitAuthError:
            if attempt == max_attempts:
                raise
            delay = min(1000 * (2 ** (attempt - 1)), 30000)
            await asyncio.sleep(delay / 1000)
    raise UrbitAuthError("Max retries exceeded")