Newer
Older
tgclient-mcp / backend / app / tg / manager.py
"""AccountManager: пул живых TelegramClient по accounts.id (один процесс).

- Клиенты создаются лениво (`get_client`), коннект кешируется; Telethon сам
  переподключается (auto_reconnect=True). `warmup()` — фоновый eager-connect
  активных аккаунтов при старте (не блокирует подъём сервиса).
- Сессии — только StringSession, сериализуются в accounts.session_data
  (aiosqlite WAL, одно соединение — см. db.py). SQLiteSession запретён:
  файл SQLite внутри контейнера конфликтует с нашей БД.
- Persist: после логина и после мутирующих вызовов (`save_session`) и
  периодический дифф-пersist (task в lifespan) — StringSession хранит
  кеш entity-хэшей, без этого get_input_entity ломается после рестарта.
- FloodWait: flood_sleep_threshold=60 — Telethon сам спит на коротких;
  длинные ловятся в тулах (429 как данные).
"""

import asyncio
from contextlib import suppress

from telethon import TelegramClient, errors
from telethon.sessions import StringSession

from app.config import get_settings
from app.db import get_db
from app.security import now_iso
from app.synapse_report import report


class AccountManager:
    def __init__(self) -> None:
        self.clients: dict[int, TelegramClient] = {}
        self._pending: dict[str, TelegramClient] = {}  # login_session_id → client

    # --- Живые аккаунты -----------------------------------------------------------

    async def get_client(self, account_id: int) -> TelegramClient:
        """Клиент по account_id: connect если нужен, результат кешируется.

        Строка accounts читается прямо из БД (aiosqlite); AuthKeyUnregistered /
        UserDeactivated → статус logged_out, клиент отбывается — DomainError.
        """
        existing = self.clients.get(account_id)
        if existing is not None and existing.is_connected():
            return existing
        row = await self._load_account(account_id)
        if row is None:
            from app.errors import DomainError

            raise DomainError(404, f"account #{account_id} not found")
        if row["status"] == "logged_out":
            raise DomainError(409, f"account #{account_id} is logged out — login again")
        client = TelegramClient(
            StringSession(row["session_data"] or ""),
            row["api_id"] or get_settings().api_id,
            row["api_hash"] or get_settings().api_hash,
            auto_reconnect=True,
            retry_delay=3,
            connection_retries=None,
            request_retries=3,
            flood_sleep_threshold=60,
            catch_up=False,
            sequential_updates=False,
        )
        await client.connect()
        authorized = await client.is_user_authorized()
        if not authorized:
            await client.disconnect()
            await self._mark_status(account_id, "logged_out")
            raise DomainError(409, f"account #{account_id} is not authorized — login again")
        self.clients[account_id] = client
        await self._touch(account_id)
        from app.tg.calls import register_call_listener
        from app.tg.events import register_listeners

        register_call_listener(client, account_id)
        if get_settings().notify_messages:
            # пуш-события сообщений → Synapse (агенты/скрипты подписаны в Synapse)
            register_listeners(client, account_id)
        return client

    async def drop(self, account_id: int) -> None:
        """Забыть клиент (disconnect): рестарт-логин, логаут, ручной disconnect."""
        client = self.clients.pop(account_id, None)
        if client is not None:
            with suppress(Exception):
                await client.disconnect()

    async def warmup(self) -> None:
        """eager-connect всех активных аккаунтов (фоновая задача; ошибки не
        блокируют старт и остальные аккаунты)."""
        db = get_db()
        cursor = await db.execute(
            "SELECT id FROM accounts WHERE status = 'active' ORDER BY id"
        )
        for row in await cursor.fetchall():
            try:
                await self.get_client(row["id"])
            except Exception as exc:  # noqa: BLE001 — один упавший не валит остальных
                print(f"warmup: account #{row['id']} failed: {exc}", flush=True)

    # --- Pending-логины ------------------------------------------------------------

    def pending_client(self, login_id: str) -> TelegramClient | None:
        return self._pending.get(login_id)

    def put_pending(self, login_id: str, client: TelegramClient) -> None:
        self._pending[login_id] = client

    async def drop_pending(self, login_id: str, *, disconnect: bool = True) -> None:
        client = self._pending.pop(login_id, None)
        if client is not None and disconnect:
            with suppress(Exception):
                await client.disconnect()

    def pending_ids(self) -> list[str]:
        return list(self._pending)

    # --- Persist сессий --------------------------------------------------------------

    async def save_session(self, account_id: int) -> None:
        """session.save() → accounts.session_data (после логина и мутаций:
        StringSession хранит entity-кеш, важно для get_input_entity)."""
        client = self.clients.get(account_id)
        if client is None:
            return
        data = client.session.save()
        db = get_db()
        await db.execute(
            "UPDATE accounts SET session_data = ?, updated_at = ? WHERE id = ?",
            (data, now_iso(), account_id),
        )
        await db.commit()

    async def persist_loop(self) -> None:
        """Периодический дифф-пersist всех живых клиентов (изменённые сессии)."""
        while True:
            try:
                for account_id, client in list(self.clients.items()):
                    stored = await self._stored_session(account_id)
                    current = client.session.save()
                    if stored != current:
                        await self.save_session(account_id)
            except Exception as exc:  # не роняем цикл
                print(f"session persist error: {exc}", flush=True)
            await asyncio.sleep(60)

    # --- Внутреннее -------------------------------------------------------------------

    async def _load_account(self, account_id: int):
        db = get_db()
        cursor = await db.execute(
            "SELECT * FROM accounts WHERE id = ?", (account_id,)
        )
        return await cursor.fetchone()

    async def _stored_session(self, account_id: int) -> str | None:
        cursor = await get_db().execute(
            "SELECT session_data FROM accounts WHERE id = ?", (account_id,)
        )
        row = await cursor.fetchone()
        return row["session_data"] if row else None

    async def _mark_status(self, account_id: int, status: str, error: str = "") -> None:
        db = get_db()
        await db.execute(
            "UPDATE accounts SET status = ?, error = ?, updated_at = ? WHERE id = ?",
            (status, error[:500], now_iso(), account_id),
        )
        await db.commit()

    async def note_auth_lost(self, account_id: int, detail: str) -> None:
        """AuthKeyUnregistered/аналог: аккаунт мёртв, клиент выкинут, Synapse
        уведомлён (ttl/dedup делает репортер)."""
        await self.drop(account_id)
        await self._mark_status(account_id, "logged_out", detail)
        db = get_db()
        cursor = await db.execute("SELECT user_id FROM accounts WHERE id = ?", (account_id,))
        row = await cursor.fetchone()
        report(
            "tg_account_auth_lost",
            {"entity": f"account-{account_id}", "user_id": row["user_id"] if row else None,
             "account_id": account_id, "detail": detail},
        )

    async def shutdown(self) -> None:
        for account_id in list(self.clients):
            self.clients[account_id].disconnect()  # fire-and-forget is fine
            self.clients.pop(account_id, None)
        for login_id in list(self._pending):
            await self.drop_pending(login_id)

    async def count(self) -> dict:
        """Счётчики для /health: active/connected/pending_logins."""
        db = get_db()
        cursor = await db.execute(
            "SELECT COUNT(*) AS n FROM accounts WHERE status = 'active'"
        )
        active = (await cursor.fetchone())["n"]
        return {"active": active, "connected": len(self.clients), "pending_logins": len(self._pending)}


async def call_with_flood_guard(manager, account_id: int, coro_fn):
    """Обёртка вызовов Telethon: длинный FloodWaitError → 429 как данные.

    Короткие (≤60 с) Telethon переспит сам (flood_sleep_threshold) и короутина
    вернётся. Длинные не ждём: агент получит причину и решит сам.
    """
    try:
        return await coro_fn()
    except errors.FloodWaitError as exc:
        report(
            "tg_flood_wait",
            {"entity": f"account-{account_id}", "account_id": account_id, "seconds": exc.seconds},
        )
        from app.errors import DomainError

        raise DomainError(429, f"telegram flood wait: {exc.seconds}s — retry later") from exc
    except errors.AuthKeyUnregisteredError as exc:
        await manager.note_auth_lost(account_id, str(exc))
        raise DomainError(409, "telegram session is dead (auth key unregistered)") from exc
    except errors.UserDeactivatedError as exc:
        await manager.note_auth_lost(account_id, str(exc))
        raise DomainError(409, "telegram user deactivated") from exc