feat(api,userbot,frontend): add accounts from the web ui and isolate per-account settings

2026-08-05 23:49:52 +02:00
parent 92fd20137e
commit 9c265af3d3
19 changed files with 917 additions and 268 deletions
@@ -1,7 +1,7 @@
 import asyncio
 import contextlib
-import json
 from collections.abc import Callable, Coroutine
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any

@@ -15,87 +15,144 @@ from userbot.modules.jobs import JobConsumer
 from utils.env import env
 from utils.jobs import enqueue
 from utils.logging import logger, setup_logging
+from utils.read.accounts import (
+    ACCOUNTS_CHANGED_CHANNEL,
+    inactive_session_names,
+    sync_account,
+)
 from utils.read.watches import WATCHES_CHANGED_CHANNEL
 from utils.storage import ContentAddressedStorage

 setup_logging()

-_UPSERT_ACCOUNT = """
-INSERT INTO accounts
-    (tg_user_id, label, phone, session_name, is_active, raw, updated_at)
-VALUES ($1, $2, $3, $4, TRUE, $5::jsonb, now())
-ON CONFLICT (tg_user_id) DO UPDATE SET
-    label = EXCLUDED.label,
-    phone = EXCLUDED.phone,
-    session_name = EXCLUDED.session_name,
-    is_active = TRUE,
-    raw = EXCLUDED.raw,
-    updated_at = now()
-RETURNING account_id
-"""
+
+@dataclass
+class RunningAccount:
+    client: PyroClient
+    consumer_task: asyncio.Task


-def _discover_sessions(sessions_dir: Path) -> list[Path]:
+def _sessions_dir() -> Path:
+    sessions_dir = Path(env.tg.sessions_dir)
    sessions_dir.mkdir(parents=True, exist_ok=True)
-    return sorted(sessions_dir.glob("*.session"))
+    return sessions_dir


-async def _sync_account(
-    pool: asyncpg.Pool, client: PyroClient, session_name: str
-) -> int | None:
-    me = client.me
-    if not me:
-        return None
-    raw = json.dumps(
-        {
-            "id": me.id,
-            "first_name": me.first_name,
-            "last_name": me.last_name,
-            "username": me.username,
-            "phone_number": me.phone_number,
-        }
-    )
-    label = " ".join(filter(None, [me.first_name, me.last_name])) or me.username
-    account_id = await pool.fetchval(
-        _UPSERT_ACCOUNT, me.id, label, me.phone_number, session_name, raw
-    )
-    logger.info(f"[green]Account synced:[/] {label} ({me.id})")
-    return account_id
+async def _cancel(task: asyncio.Task) -> None:
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await task


-async def _setup_capture(
-    pool: asyncpg.Pool,
-    client: PyroClient,
-    account_id: int,
-    storage: ContentAddressedStorage,
-) -> None:
-    client.capture = await build_capture_context(client, pool, storage, account_id)
-    logger.info("[green]Capture context ready.[/]")
-
-
-async def _enqueue_sync_dialogs(pool: asyncpg.Pool, account_id: int) -> None:
+async def _enqueue_once(pool: asyncpg.Pool, account_id: int, kind: str) -> None:
    existing = await pool.fetchval(
-        "SELECT 1 FROM jobs WHERE account_id = $1 AND kind = 'sync_dialogs' "
+        "SELECT 1 FROM jobs WHERE account_id = $1 AND kind = $2 "
        "AND status IN ('pending', 'running') LIMIT 1",
        account_id,
+        kind,
    )
    if existing is None:
-        await enqueue(pool, account_id, "sync_dialogs", {})
-        logger.info("[green]Queued sync_dialogs.[/]")
+        await enqueue(pool, account_id, kind, {})
+        logger.info(f"[green]Queued {kind}.[/]")
+
+
+class AccountRegistry:
+    def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None:
+        self._pool = pool
+        self._storage = storage
+        self._running: dict[str, RunningAccount] = {}
+        self._lock = asyncio.Lock()
+
+    @property
+    def clients(self) -> list[PyroClient]:
+        return [account.client for account in self._running.values()]
+
+    async def sync(self) -> None:
+        async with self._lock:
+            inactive = await inactive_session_names(self._pool)
+            present: set[str] = set()
+            for path in sorted(_sessions_dir().glob("*.session")):
+                if path.stem in inactive:
+                    await self._log_out(path)
+                    continue
+                present.add(path.stem)
+                if path.stem not in self._running:
+                    await self._start(path)
+            for session_name in set(self._running) - present:
+                await self._stop(session_name)
+
+    async def close(self) -> None:
+        for session_name in list(self._running):
+            await self._stop(session_name)
+
+    async def _start(self, path: Path) -> None:
+        session_name = path.stem
+        client = PyroClient(session_name, workdir=str(path.parent))
+        try:
+            await client.start()
+            me = client.me
+            if me is None:
+                msg = f"session {session_name} is not authorized"
+                raise RuntimeError(msg)
+            account_id = await sync_account(self._pool, me, session_name)
+            client.capture = await build_capture_context(
+                client, self._pool, self._storage, account_id
+            )
+        except Exception:
+            logger.exception(f"[red]Failed to start session:[/] {session_name}")
+            with contextlib.suppress(Exception):
+                await client.stop()
+            return
+        consumer = JobConsumer(client, self._pool, account_id)
+        self._running[session_name] = RunningAccount(
+            client, asyncio.create_task(consumer.run())
+        )
+        logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})")
+        await _enqueue_once(self._pool, account_id, "sync_dialogs")
+        await _enqueue_once(self._pool, account_id, "sync_contacts")
+
+    async def _stop(self, session_name: str) -> None:
+        account = self._running.pop(session_name, None)
+        if account is None:
+            return
+        await _cancel(account.consumer_task)
+        with contextlib.suppress(Exception):
+            await account.client.stop()
+        logger.info(f"[yellow]Client stopped:[/] {session_name}")
+
+    async def _log_out(self, path: Path) -> None:
+        account = self._running.pop(path.stem, None)
+        if account is not None:
+            await _cancel(account.consumer_task)
+            client = account.client
+        else:
+            client = PyroClient(
+                path.stem, workdir=str(path.parent), load_handlers=False
+            )
+        with contextlib.suppress(Exception):
+            if account is None:
+                await client.start()
+            await client.log_out()
+        with contextlib.suppress(Exception):
+            await client.stop()
+        path.unlink(missing_ok=True)  # noqa: ASYNC240
+        logger.info(f"[yellow]Account logged out:[/] {path.stem}")


 async def _listen_changes(
-    clients: list[PyroClient], tasks: set[asyncio.Task]
+    registry: AccountRegistry, tasks: set[asyncio.Task]
 ) -> asyncpg.Connection:
+    def spawn(coro: Coroutine[Any, Any, None]) -> None:
+        task = asyncio.create_task(coro)
+        tasks.add(task)
+        task.add_done_callback(tasks.discard)
+
    def reload(
        make_coro: Callable[[CaptureContext], Coroutine[Any, Any, None]],
    ) -> None:
-        for client in clients:
-            if client.capture is None:
-                continue
-            task = asyncio.create_task(make_coro(client.capture))
-            tasks.add(task)
-            task.add_done_callback(tasks.discard)
+        for client in registry.clients:
+            if client.capture is not None:
+                spawn(make_coro(client.capture))

    def on_policy(
        _conn: asyncpg.Connection, _pid: int, _channel: str, _payload: str
@@ -107,9 +164,15 @@ async def _listen_changes(
    ) -> None:
        reload(lambda capture: capture.watches.refresh())

+    def on_accounts(
+        _conn: asyncpg.Connection, _pid: int, _channel: str, _payload: str
+    ) -> None:
+        spawn(registry.sync())
+
    conn = await asyncpg.connect(dsn=env.db.connection_url)
    await conn.add_listener("policy_changed", on_policy)
    await conn.add_listener(WATCHES_CHANGED_CHANNEL, on_watch)
+    await conn.add_listener(ACCOUNTS_CHANGED_CHANNEL, on_accounts)
    return conn


@@ -117,53 +180,21 @@ async def runner() -> None:
    pool = await container.get(asyncpg.Pool)
    storage = await container.get(ContentAddressedStorage)

-    sessions_dir = Path(env.tg.sessions_dir)
-    session_files = _discover_sessions(sessions_dir)
-
-    if not session_files:
-        logger.warning(
-            f"[yellow]No .session files in {sessions_dir}/. "
-            f"Log in first, then restart userbot.[/]"
-        )
-
-    clients: list[PyroClient] = []
-    reload_tasks: set[asyncio.Task] = set()
-    consumer_tasks: list[asyncio.Task] = []
+    registry = AccountRegistry(pool, storage)
+    tasks: set[asyncio.Task] = set()
    listen_conn: asyncpg.Connection | None = None
    try:
-        for session_path in session_files:
-            session_name = session_path.stem
-            client = PyroClient(session_name, workdir=str(sessions_dir))
-            await client.start()
-            clients.append(client)
-            logger.info(
-                f"[green]Client started:[/] "
-                f"{client.me.full_name if client.me else 'unknown'} "
-                f"{client.me.id if client.me else 'unknown'}"
-            )
-            account_id = await _sync_account(pool, client, session_name)
-            if account_id is not None:
-                await _setup_capture(pool, client, account_id, storage)
-                consumer = JobConsumer(client, pool, account_id)
-                consumer_tasks.append(asyncio.create_task(consumer.run()))
-                await _enqueue_sync_dialogs(pool, account_id)
-
-        if clients:
-            listen_conn = await _listen_changes(clients, reload_tasks)
-            logger.info("[green]Userbot running.[/]")
-            await asyncio.Event().wait()
+        await registry.sync()
+        if not registry.clients:
+            logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]")
+        listen_conn = await _listen_changes(registry, tasks)
+        logger.info("[green]Userbot running.[/]")
+        await asyncio.Event().wait()
    finally:
-        for task in consumer_tasks:
-            task.cancel()
-        for task in consumer_tasks:
-            with contextlib.suppress(asyncio.CancelledError):
-                await task
        if listen_conn is not None:
            with contextlib.suppress(Exception):
                await listen_conn.close()
-        for client in clients:
-            with contextlib.suppress(Exception):
-                await client.stop()
+        await registry.close()
        await container.close()