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

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}")