feat(userbot,api,frontend): backfill stories and unbreak all-stories list

backend/src/api/routers/backfill.py

@@ -19,6 +19,11 @@ class BackfillRequest(BaseModel):
     full: bool = False
 
 
+class StoriesBackfillRequest(BaseModel):
+    account_id: int
+    peer_id: int
+
+
 class FetchMediaRequest(BaseModel):
     account_id: int
     chat_id: int
@@ -76,6 +81,16 @@ async def enqueue_backfill(
     return EnqueueResponse(job_id=job_id)
 
 
+@router.post("/stories/backfill", status_code=201)
+async def enqueue_stories_backfill(
+    pool: FromDishka[asyncpg.Pool], body: StoriesBackfillRequest
+) -> EnqueueResponse:
+    job_id = await enqueue(
+        pool, body.account_id, "backfill_stories", {"peer_id": body.peer_id}
+    )
+    return EnqueueResponse(job_id=job_id)
+
+
 @router.post("/media/fetch", status_code=201)
 async def enqueue_fetch_media(
     pool: FromDishka[asyncpg.Pool], body: FetchMediaRequest

backend/src/userbot/handlers/stories.py

@@ -1,52 +1,14 @@
-from io import BytesIO
-
 from pyrogram.types import Story
 
 from userbot import PyroClient
-from userbot.modules.stories import repository
-
-
-def _peer_id(story: Story) -> int:
-    if story.chat is not None:
-        return story.chat.id or 0
-    if story.from_user is not None:
-        return story.from_user.id or 0
-    return 0
+from userbot.modules.stories.service import save_story
 
 
 @PyroClient.on_story()
 async def on_story(client: PyroClient, story: Story) -> None:
-    ctx = client.capture
-    if ctx is None:
+    if client.capture is None:
         return
-    media_kind = story.media.name.lower() if story.media else None
-    storage_key: str | None = None
-    file_size: int | None = None
-    downloaded = False
-    if not story.deleted and story.media is not None:
-        buffer = await client.download_media(story, in_memory=True)
-        if isinstance(buffer, BytesIO):
-            data = buffer.getvalue()
-            storage_key = ctx.storage.put(data)
-            file_size = len(data)
-            downloaded = True
-    await repository.upsert_story(
-        ctx.pool,
-        ctx.account_id,
-        _peer_id(story),
-        story.id,
-        story.date,
-        story.expire_date,
-        story.caption,
-        media_kind,
-        storage_key,
-        file_size,
-        story.views,
-        str(story.raw),
-        pinned=bool(story.pinned),
-        deleted=bool(story.deleted),
-        downloaded=downloaded,
-    )
+    await save_story(client, client.capture, story)
 
 
 handlers = on_story.handlers

backend/src/userbot/modules/jobs/handlers/__init__.py

@@ -1,5 +1,6 @@
 from userbot.modules.jobs.handlers import (
     backfill,
+    backfill_stories,
     enrich_chat,
     fetch_avatar,
     fetch_custom_emoji,
@@ -12,6 +13,7 @@ from userbot.modules.jobs.handlers import (
 
 __all__ = [
     "backfill",
+    "backfill_stories",
     "enrich_chat",
     "fetch_avatar",
     "fetch_custom_emoji",

backend/src/userbot/modules/jobs/handlers/backfill_stories.py

@@ -0,0 +1,70 @@
+from collections.abc import AsyncIterator, Callable
+
+from pyrogram import Client
+from pyrogram.errors import FloodPremiumWait, FloodWait, RPCError
+from pyrogram.types import Story
+
+from userbot.modules.capture.context import CaptureContext
+from userbot.modules.jobs.context import JobContext
+from userbot.modules.jobs.registry import register
+from userbot.modules.stories.service import save_story
+
+SAVE_EVERY = 10
+
+StorySource = Callable[[], AsyncIterator[Story]]
+
+
+def _sources(client: Client, peer_id: int, *, own: bool) -> dict[str, StorySource]:
+    sources: dict[str, StorySource] = {
+        "active": lambda: client.get_chat_stories(peer_id),
+        "pinned": lambda: client.get_pinned_stories(peer_id),
+    }
+    if own:
+        sources["archived"] = lambda: client.get_archived_stories(peer_id)
+    return sources
+
+
+async def _drain(
+    ctx: JobContext, capture: CaptureContext, name: str, source: StorySource
+) -> int:
+    client = ctx.client
+    if client is None:
+        return 0
+    saved = 0
+    async for story in source():
+        try:
+            await save_story(client, capture, story)
+        except (FloodWait, FloodPremiumWait):
+            raise
+        except RPCError:
+            continue
+        saved += 1
+        if saved % SAVE_EVERY == 0:
+            await ctx.report_progress({"saved": saved, "source": name})
+            if await ctx.is_canceled():
+                break
+    return saved
+
+
+@register("backfill_stories")
+async def backfill_stories(ctx: JobContext) -> None:
+    client = ctx.client
+    if client is None:
+        return
+    capture = getattr(client, "capture", None)
+    if capture is None:
+        return
+    peer_id = ctx.job.params["peer_id"]
+    own = client.me is not None and client.me.id == peer_id
+    saved = 0
+    errors: dict[str, str] = {}
+    for name, source in _sources(client, peer_id, own=own).items():
+        try:
+            saved += await _drain(ctx, capture, name, source)
+        except (FloodWait, FloodPremiumWait):
+            raise
+        except RPCError as exc:
+            errors[name] = type(exc).__name__
+        if await ctx.is_canceled():
+            break
+    await ctx.report_progress({"saved": saved, "done": True, "errors": errors})

backend/src/userbot/modules/stories/repository.py

@@ -22,6 +22,20 @@ ON CONFLICT (account_id, peer_id, story_id) DO UPDATE SET
 """
 
 
+async def is_downloaded(
+    pool: asyncpg.Pool, account_id: int, peer_id: int, story_id: int
+) -> bool:
+    return bool(
+        await pool.fetchval(
+            "SELECT downloaded FROM stories "
+            "WHERE account_id = $1 AND peer_id = $2 AND story_id = $3",
+            account_id,
+            peer_id,
+            story_id,
+        )
+    )
+
+
 async def upsert_story(  # noqa: PLR0913
     pool: asyncpg.Pool,
     account_id: int,

backend/src/userbot/modules/stories/service.py

@@ -0,0 +1,49 @@
+from io import BytesIO
+
+from pyrogram import Client
+from pyrogram.types import Story
+
+from userbot.modules.capture.context import CaptureContext
+from userbot.modules.stories import repository
+
+
+def story_peer_id(story: Story) -> int:
+    if story.chat is not None:
+        return story.chat.id or 0
+    if story.from_user is not None:
+        return story.from_user.id or 0
+    return 0
+
+
+async def save_story(client: Client, capture: CaptureContext, story: Story) -> None:
+    peer_id = story_peer_id(story)
+    storage_key: str | None = None
+    file_size: int | None = None
+    downloaded = False
+    stored = await repository.is_downloaded(
+        capture.pool, capture.account_id, peer_id, story.id
+    )
+    if not (stored or story.deleted or story.media is None):
+        buffer = await client.download_media(story, in_memory=True)
+        if isinstance(buffer, BytesIO):
+            data = buffer.getvalue()
+            storage_key = capture.storage.put(data)
+            file_size = len(data)
+            downloaded = True
+    await repository.upsert_story(
+        capture.pool,
+        capture.account_id,
+        peer_id,
+        story.id,
+        story.date,
+        story.expire_date,
+        story.caption,
+        story.media.name.lower() if story.media else None,
+        storage_key,
+        file_size,
+        story.views,
+        str(story.raw),
+        pinned=bool(story.pinned),
+        deleted=bool(story.deleted),
+        downloaded=downloaded,
+    )