feat(userbot): after a restart the userbot catches up on what it missed while down, from a saved cursor and within bounds

backend/src/userbot/runner.py

@@ -11,6 +11,7 @@ import uvloop
 
 from dependencies.container import container
 from userbot.modules.capture import CaptureContext, build_capture_context
+from userbot.modules.catchup import SavedState, catch_up, read_saved_state
 from userbot.modules.client import PyroClient
 from userbot.modules.jobs import JobConsumer
 from utils.env import env
@@ -35,7 +36,7 @@ SESSION_RETRY_INTERVAL = 60
 @dataclass
 class RunningAccount:
     client: PyroClient
-    consumer_task: asyncio.Task
+    tasks: list[asyncio.Task]
     device_model: str | None
 
 
@@ -56,6 +57,7 @@ class AccountRegistry:
         self._pool = pool
         self._storage = storage
         self._running: dict[str, RunningAccount] = {}
+        self._saved: dict[str, SavedState] = {}
         self._lock = asyncio.Lock()
 
     @property
@@ -88,6 +90,8 @@ class AccountRegistry:
 
     async def _start(self, path: Path, device_model: str | None) -> None:
         session_name = path.stem
+        if session_name not in self._saved:
+            self._saved[session_name] = await asyncio.to_thread(read_saved_state, path)
         client = PyroClient(
             session_name, workdir=str(path.parent), device_model=device_model
         )
@@ -110,9 +114,12 @@ class AccountRegistry:
                 await client.stop()
             return
         consumer = JobConsumer(client, self._pool, account_id)
-        self._running[session_name] = RunningAccount(
-            client, asyncio.create_task(consumer.run()), device_model
-        )
+        saved = self._saved.pop(session_name)
+        tasks = [
+            asyncio.create_task(consumer.run()),
+            asyncio.create_task(catch_up(client, saved)),
+        ]
+        self._running[session_name] = RunningAccount(client, tasks, device_model)
         logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})")
         await enqueue(self._pool, account_id, "sync_dialogs", {})
         await enqueue(self._pool, account_id, "sync_contacts", {})
@@ -121,7 +128,8 @@ class AccountRegistry:
         account = self._running.pop(session_name, None)
         if account is None:
             return
-        await _cancel(account.consumer_task)
+        for task in account.tasks:
+            await _cancel(task)
         with contextlib.suppress(Exception):
             await account.client.stop()
         logger.info(f"[yellow]Client stopped:[/] {session_name}")
@@ -129,7 +137,8 @@ class AccountRegistry:
     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)
+            for task in account.tasks:
+                await _cancel(task)
             client = account.client
         else:
             client = PyroClient(