feat(watcher,api,frontend): watched bots become accounts and drain the update backlog

backend/src/watcher/backfill.py

@@ -1,71 +0,0 @@
-import asyncio
-
-import asyncpg
-from pyrogram import Client
-from pyrogram.errors import FloodWait, RPCError
-
-from utils.env import env
-from utils.logging import logger
-from utils.read import bots
-from watcher.capture import SOURCE_HISTORY, Row, message_row, store_many
-
-BATCH = 200
-MAX_FLOOD_WAIT = 300
-HISTORY_UNAVAILABLE = "История недоступна: {}"
-
-
-async def _flush(pool: asyncpg.Pool, rows: list[Row]) -> int:
-    await store_many(pool, rows)
-    stored = len(rows)
-    rows.clear()
-    return stored
-
-
-async def _chat_history(
-    client: Client, pool: asyncpg.Pool, bot_id: int, chat_id: int
-) -> int:
-    rows: list[Row] = []
-    stored = 0
-    async for message in client.get_chat_history(
-        chat_id, limit=env.bots.backfill_per_chat
-    ):
-        if message.empty or message.chat is None or message.date is None:
-            continue
-        rows.append(message_row(bot_id, message, source=SOURCE_HISTORY))
-        if len(rows) >= BATCH:
-            stored += await _flush(pool, rows)
-    return stored + await _flush(pool, rows)
-
-
-async def backfill(client: Client, pool: asyncpg.Pool, bot_id: int) -> None:
-    try:
-        dialogs = [
-            dialog.chat.id
-            async for dialog in client.get_dialogs(limit=env.bots.backfill_chats)
-            if dialog.chat is not None and dialog.chat.id is not None
-        ]
-    except RPCError as exc:
-        logger.warning(f"[yellow]Bot {bot_id}: dialogs unavailable:[/] {exc}")
-        await bots.set_status(
-            pool, bot_id, bots.STATUS_OK, detail=HISTORY_UNAVAILABLE.format(exc)
-        )
-        return
-
-    stored = 0
-    for chat_id in dialogs:
-        try:
-            stored += await _chat_history(client, pool, bot_id, chat_id)
-        except FloodWait as exc:
-            wait = exc.value if isinstance(exc.value, int) else MAX_FLOOD_WAIT + 1
-            if wait > MAX_FLOOD_WAIT:
-                logger.warning(f"[yellow]Bot {bot_id}: flood wait {wait}s, stop[/]")
-                break
-            await asyncio.sleep(wait)
-        except RPCError as exc:
-            logger.warning(f"[yellow]Bot {bot_id}: chat {chat_id} skipped:[/] {exc}")
-
-    await bots.mark_backfilled(pool, bot_id)
-    logger.info(
-        f"[green]Bot {bot_id}: backfilled[/] {stored} messages "
-        f"from {len(dialogs)} chats"
-    )

backend/src/watcher/capture.py

@@ -1,146 +0,0 @@
-from datetime import UTC, datetime
-
-import asyncpg
-from pyrogram.types import CallbackQuery, ChatMemberUpdated, Message, User
-
-from userbot.modules.raw_json import raw_json
-from utils.read import bots
-
-SOURCE_HISTORY = "history"
-SOURCE_LIVE = "live"
-
-KIND_MESSAGE = "message"
-KIND_CALLBACK = "callback_query"
-KIND_MEMBER = "member"
-
-MEDIA_KINDS = (
-    "photo",
-    "video",
-    "animation",
-    "audio",
-    "voice",
-    "video_note",
-    "document",
-    "sticker",
-    "contact",
-    "location",
-    "venue",
-    "poll",
-    "dice",
-    "story",
-    "game",
-)
-
-Row = tuple
-
-
-def _media_kind(message: Message) -> str | None:
-    return next(
-        (kind for kind in MEDIA_KINDS if getattr(message, kind, None) is not None), None
-    )
-
-
-def _full_name(user: User | None) -> str | None:
-    if user is None:
-        return None
-    parts = [user.first_name, user.last_name]
-    return " ".join(part for part in parts if part) or None
-
-
-def _chat_title(message: Message) -> str | None:
-    chat = message.chat
-    if chat is None:
-        return None
-    parts = [chat.first_name, chat.last_name]
-    return chat.title or " ".join(part for part in parts if part) or None
-
-
-def _chat_type(message: Message) -> str | None:
-    kind = getattr(message.chat, "type", None)
-    return kind.value if kind is not None else None
-
-
-def message_row(bot_id: int, message: Message, *, source: str) -> Row:
-    user = message.from_user
-    chat = message.chat
-    sender = user.id if user is not None else None
-    if sender is None and message.sender_chat is not None:
-        sender = message.sender_chat.id
-    return (
-        bot_id,
-        KIND_MESSAGE,
-        chat.id if chat is not None else 0,
-        message.id,
-        message.date,
-        source,
-        bool(message.outgoing),
-        sender,
-        user.username if user is not None else None,
-        _full_name(user),
-        user.is_bot if user is not None else None,
-        _chat_type(message),
-        _chat_title(message),
-        message.text or message.caption,
-        _media_kind(message),
-        message.edit_date,
-        raw_json(message),
-    )
-
-
-def callback_row(bot_id: int, query: CallbackQuery) -> Row:
-    user = query.from_user
-    message = query.message
-    chat = message.chat if message is not None else None
-    return (
-        bot_id,
-        KIND_CALLBACK,
-        chat.id if chat is not None else 0,
-        message.id if message is not None else 0,
-        datetime.now(UTC),
-        SOURCE_LIVE,
-        False,
-        user.id if user is not None else None,
-        user.username if user is not None else None,
-        _full_name(user),
-        user.is_bot if user is not None else None,
-        chat.type.value if chat is not None and chat.type is not None else None,
-        chat.title if chat is not None else None,
-        query.data,
-        None,
-        None,
-        raw_json(query),
-    )
-
-
-def member_row(bot_id: int, update: ChatMemberUpdated) -> Row:
-    member = update.new_chat_member
-    user = member.user if member is not None else None
-    chat = update.chat
-    status = member.status.value if member is not None else None
-    return (
-        bot_id,
-        KIND_MEMBER,
-        chat.id if chat is not None else 0,
-        user.id if user is not None else 0,
-        update.date or datetime.now(UTC),
-        SOURCE_LIVE,
-        False,
-        update.from_user.id if update.from_user is not None else None,
-        update.from_user.username if update.from_user is not None else None,
-        _full_name(update.from_user),
-        update.from_user.is_bot if update.from_user is not None else None,
-        chat.type.value if chat is not None and chat.type is not None else None,
-        chat.title if chat is not None else None,
-        status,
-        None,
-        None,
-        raw_json(update),
-    )
-
-
-async def store(pool: asyncpg.Pool, row: Row) -> None:
-    await bots.store_update(pool, row)
-
-
-async def store_many(pool: asyncpg.Pool, rows: list[Row]) -> None:
-    await bots.store_updates(pool, rows)

backend/src/watcher/context.py

@@ -0,0 +1,33 @@
+"""A capture context for a bot session.
+
+Same context the userbot capture path expects, minus the two caches a bot
+may not fill: ``messages.getDialogFilters`` and ``contacts.getContacts`` are
+both refused for bots, so folders and contacts simply stay empty and policy
+falls back to the per-account defaults.
+"""
+
+import asyncpg
+from pyrogram import Client
+
+from userbot.modules.capture.context import CaptureContext
+from userbot.modules.contacts import ContactCache
+from userbot.modules.folders import FolderCache
+from utils.storage import ContentAddressedStorage
+
+
+async def build_bot_capture_context(
+    client: Client,
+    pool: asyncpg.Pool,
+    storage: ContentAddressedStorage,
+    account_id: int,
+) -> CaptureContext:
+    ctx = CaptureContext(
+        account_id,
+        pool,
+        storage,
+        FolderCache(client, pool, account_id),
+        ContactCache(client),
+    )
+    await ctx.reload_policies()
+    await ctx.watches.refresh()
+    return ctx

backend/src/watcher/drain.py

@@ -0,0 +1,96 @@
+"""Pull the updates Telegram still holds for a bot that nobody consumed.
+
+A bot's undelivered updates outlive the session they were meant for: the
+same queue Bot API ``getUpdates`` drains is reachable over MTProto through
+``updates.getDifference``. A fresh login normally starts from
+``updates.getState``, which throws that queue away - this walks it instead.
+"""
+
+from pyrogram import raw, types
+from pyrogram.errors import RPCError
+
+from userbot import PyroClient
+from userbot.modules.capture import capture_message
+from userbot.modules.capture.chat_meta import meta_from_chat
+from userbot.modules.capture.context import CaptureContext
+from utils.logging import logger
+
+EARLIEST_PTS = 1
+EARLIEST_QTS = 1
+MAX_SLICES = 50
+
+
+async def _store(
+    client: PyroClient, ctx: CaptureContext, message: types.Message
+) -> int:
+    if message.empty or message.chat is None or message.date is None:
+        return 0
+    meta = meta_from_chat(message.chat, ctx.contacts.ids)
+    toggles = ctx.resolve(meta)
+    if not toggles.messages:
+        return 0
+    await capture_message(client, message, ctx, toggles)
+    return 1
+
+
+async def _parse_all(
+    client: PyroClient, new_messages: list, users: list, chats: list
+) -> list[types.Message]:
+    user_map = {u.id: u for u in users}
+    chat_map = {c.id: c for c in chats}
+    parsed = []
+    for raw_message in new_messages:
+        if isinstance(raw_message, raw.types.MessageEmpty):
+            continue
+        parsed.append(
+            await types.Message._parse(client, raw_message, user_map, chat_map)  # noqa: SLF001
+        )
+    return parsed
+
+
+async def drain_backlog(client: PyroClient, account_id: int) -> int:
+    """Walk the bot's pending update queue into the archive. Returns stored count."""
+    ctx = client.capture
+    if ctx is None:
+        return 0
+    pts, qts, date = EARLIEST_PTS, EARLIEST_QTS, 0
+    stored = 0
+    for _ in range(MAX_SLICES):
+        try:
+            diff = await client.invoke(
+                raw.functions.updates.GetDifference(pts=pts, date=date, qts=qts)
+            )
+        except RPCError as exc:
+            logger.warning(
+                f"[yellow]Account {account_id}: getDifference failed:[/] {exc}"
+            )
+            return stored
+
+        if isinstance(diff, raw.types.updates.DifferenceEmpty):
+            logger.info(f"Account {account_id}: backlog empty at pts={pts}")
+            return stored
+        if isinstance(diff, raw.types.updates.DifferenceTooLong):
+            logger.info(
+                f"Account {account_id}: backlog too long, server pts={diff.pts}; "
+                f"nothing older is retrievable"
+            )
+            return stored
+
+        for message in await _parse_all(
+            client, diff.new_messages, diff.users, diff.chats
+        ):
+            stored += await _store(client, ctx, message)
+
+        if isinstance(diff, raw.types.updates.Difference):
+            logger.info(
+                f"[green]Account {account_id}: backlog drained,[/] {stored} messages"
+            )
+            return stored
+
+        pts, qts, date = (
+            diff.intermediate_state.pts,
+            diff.intermediate_state.qts,
+            diff.intermediate_state.date,
+        )
+    logger.warning(f"[yellow]Account {account_id}: backlog walk hit the slice cap[/]")
+    return stored

backend/src/watcher/registry.py

@@ -3,31 +3,37 @@ import contextlib
 from dataclasses import dataclass
 
 import asyncpg
-from pyrogram import Client
+from pyrogram import Client, enums
 from pyrogram.errors import RPCError
 from pyrogram.handlers import (
-    CallbackQueryHandler,
     ChatMemberUpdatedHandler,
     EditedMessageHandler,
     MessageHandler,
 )
 
+from userbot import PyroClient
+from userbot.modules.capture import capture_message
+from userbot.modules.capture.chat_meta import meta_from_chat
+from userbot.modules.jobs import JobConsumer
+from utils.events import notify_bg_event
+from utils.jobs import enqueue
 from utils.logging import logger
-from utils.read import bots
-from utils.telegram import API_HASH, API_ID
+from utils.read import accounts, bots
+from utils.storage import ContentAddressedStorage
 from watcher import tokens
-from watcher.backfill import backfill
-from watcher.capture import SOURCE_LIVE, callback_row, member_row, message_row, store
+from watcher.context import build_bot_capture_context
+from watcher.drain import drain_backlog
 
-DEVICE_MODEL = "Beavergram Watcher"
+DEVICE_MODEL = "Beavergram Bot"
 TOKEN_MISSING = "Токен не найден на диске"  # noqa: S105
 INVALID_TOKEN_ERRORS = ("ACCESS_TOKEN_INVALID", "ACCESS_TOKEN_EXPIRED")
 
 
 @dataclass
 class RunningBot:
-    client: Client
-    backfill_task: asyncio.Task | None
+    client: PyroClient
+    account_id: int
+    tasks: list[asyncio.Task]
 
 
 async def _cancel(task: asyncio.Task) -> None:
@@ -36,22 +42,40 @@ async def _cancel(task: asyncio.Task) -> None:
         await task
 
 
-def _handlers(pool: asyncpg.Pool, bot_id: int) -> tuple:
-    async def on_message(_client: Client, message) -> None:  # noqa: ANN001
-        if message.empty or message.chat is None or message.date is None:
+def _handlers(pool: asyncpg.Pool) -> tuple:
+    async def on_message(client: Client, message) -> None:  # noqa: ANN001
+        ctx = getattr(client, "capture", None)
+        if ctx is None or message.empty or message.chat is None or message.date is None:
             return
-        await store(pool, message_row(bot_id, message, source=SOURCE_LIVE))
+        meta = meta_from_chat(message.chat, ctx.contacts.ids)
+        await ctx.watches.on_text(
+            meta.chat_id, message.id, message.text or message.caption
+        )
+        toggles = ctx.resolve(meta)
+        if not toggles.messages:
+            return
+        await capture_message(client, message, ctx, toggles)
+        await notify_bg_event(
+            ctx.pool,
+            "message",
+            ctx.account_id,
+            chat_id=meta.chat_id,
+            message_id=message.id,
+        )
+        if message.chat.type is not enums.ChatType.PRIVATE:
+            await enqueue(
+                pool, ctx.account_id, "enrich_chat", {"chat_id": meta.chat_id}
+            )
 
-    async def on_callback(_client: Client, query) -> None:  # noqa: ANN001
-        await store(pool, callback_row(bot_id, query))
-
-    async def on_member(_client: Client, update) -> None:  # noqa: ANN001
-        await store(pool, member_row(bot_id, update))
+    async def on_member(client: Client, update) -> None:  # noqa: ANN001
+        ctx = getattr(client, "capture", None)
+        if ctx is None or update.chat is None:
+            return
+        await enqueue(pool, ctx.account_id, "enrich_chat", {"chat_id": update.chat.id})
 
     return (
         MessageHandler(on_message),
         EditedMessageHandler(on_message),
-        CallbackQueryHandler(on_callback),
         ChatMemberUpdatedHandler(on_member),
     )
 
@@ -61,8 +85,9 @@ def _is_invalid_token(exc: RPCError) -> bool:
 
 
 class BotRegistry:
-    def __init__(self, pool: asyncpg.Pool) -> None:
+    def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None:
         self._pool = pool
+        self._storage = storage
         self._running: dict[int, RunningBot] = {}
         self._lock = asyncio.Lock()
 
@@ -91,6 +116,16 @@ class BotRegistry:
             await bots.register_bot(self._pool, bot_id)
             logger.info(f"[green]Bot token adopted from disk:[/] {bot_id}")
 
+    def _client(self, bot_id: int, token: str) -> PyroClient:
+        return PyroClient(
+            str(bot_id),
+            workdir=str(tokens.sessions_dir()),
+            device_model=DEVICE_MODEL,
+            load_handlers=False,
+            bot_token=token,
+            skip_updates=False,
+        )
+
     async def _start(self, bot_id: int) -> None:
         token = tokens.load(bot_id)
         if token is None:
@@ -98,15 +133,8 @@ class BotRegistry:
                 self._pool, bot_id, bots.STATUS_UNAUTHORIZED, detail=TOKEN_MISSING
             )
             return
-        client = Client(
-            str(bot_id),
-            api_id=API_ID,
-            api_hash=API_HASH,
-            bot_token=token,
-            workdir=str(tokens.sessions_dir()),
-            device_model=DEVICE_MODEL,
-        )
-        for handler in _handlers(self._pool, bot_id):
+        client = self._client(bot_id, token)
+        for handler in _handlers(self._pool):
             client.add_handler(handler)
         try:
             await client.start()
@@ -134,23 +162,39 @@ class BotRegistry:
                 await client.stop()
             return
 
+        account_id = await accounts.sync_account(
+            self._pool, me, str(bot_id), is_bot=True
+        )
+        client.capture = await build_bot_capture_context(
+            client, self._pool, self._storage, account_id
+        )
+        await bots.link_account(self._pool, bot_id, account_id)
         await bots.set_status(
             self._pool, bot_id, bots.STATUS_OK, username=me.username, title=me.full_name
         )
-        task = (
-            asyncio.create_task(backfill(client, self._pool, bot_id))
-            if await bots.needs_backfill(self._pool, bot_id)
-            else None
+        await accounts.notify_accounts_changed(self._pool)
+
+        consumer = JobConsumer(client, self._pool, account_id)
+        tasks = [asyncio.create_task(consumer.run())]
+        if await bots.needs_backfill(self._pool, bot_id):
+            tasks.append(asyncio.create_task(self._backlog(client, bot_id, account_id)))
+        self._running[bot_id] = RunningBot(client, account_id, tasks)
+        logger.info(
+            f"[green]Bot account {account_id} watched:[/] @{me.username} ({bot_id})"
         )
-        self._running[bot_id] = RunningBot(client, task)
-        logger.info(f"[green]Bot watched:[/] @{me.username} ({bot_id})")
+
+    async def _backlog(self, client: PyroClient, bot_id: int, account_id: int) -> None:
+        stored = await drain_backlog(client, account_id)
+        await bots.mark_backfilled(self._pool, bot_id)
+        if stored:
+            await accounts.notify_accounts_changed(self._pool)
 
     async def _stop(self, bot_id: int) -> None:
         running = self._running.pop(bot_id, None)
         if running is None:
             return
-        if running.backfill_task is not None:
-            await _cancel(running.backfill_task)
+        for task in running.tasks:
+            await _cancel(task)
         with contextlib.suppress(Exception):
             await running.client.stop()
         logger.info(f"[yellow]Bot unwatched:[/] {bot_id}")

backend/src/watcher/runner.py

@@ -9,6 +9,7 @@ from dependencies.container import container
 from utils.env import env
 from utils.logging import logger, setup_logging
 from utils.read.bots import BOTS_CHANGED_CHANNEL
+from utils.storage import ContentAddressedStorage
 from watcher.registry import BotRegistry
 
 setup_logging()
@@ -40,7 +41,8 @@ async def _listen_changes(
 
 async def runner() -> None:
     pool = await container.get(asyncpg.Pool)
-    registry = BotRegistry(pool)
+    storage = await container.get(ContentAddressedStorage)
+    registry = BotRegistry(pool, storage)
     tasks: set[asyncio.Task] = set()
     listen_conn: asyncpg.Connection | None = None
     stop = asyncio.Event()