fix(userbot): drop stale media sessions and survive download timeouts

2026-09-01 21:05:48 +02:00
parent cf277c56cf
commit 9c12628304
6 changed files with 83 additions and 24 deletions
@@ -0,0 +1,56 @@
+import asyncio
+import contextlib
+from io import BytesIO
+
+from pyrogram import Client
+from pyrogram.types import (
+    Animation,
+    Audio,
+    Document,
+    Message,
+    Photo,
+    Sticker,
+    Story,
+    Video,
+    VideoNote,
+    Voice,
+)
+
+from utils.logging import logger
+
+type Downloadable = (
+    str
+    | Message
+    | Story
+    | Audio
+    | Document
+    | Photo
+    | Sticker
+    | Animation
+    | Video
+    | Voice
+    | VideoNote
+)
+
+STOP_TIMEOUT = 5.0
+
+_reset_lock = asyncio.Lock()
+
+
+async def _reset_media_sessions(client: Client) -> None:
+    async with _reset_lock:
+        sessions = list(client.media_sessions.values())
+        client.media_sessions.clear()
+        for session in sessions:
+            with contextlib.suppress(Exception):
+                await asyncio.wait_for(session.stop(), STOP_TIMEOUT)
+        logger.warning(f"[yellow]Dropped {len(sessions)} stale media session(s).[/]")
+
+
+async def download_bytes(client: Client, target: Downloadable) -> bytes | None:
+    try:
+        buffer = await client.download_media(target, in_memory=True)
+    except TimeoutError:
+        await _reset_media_sessions(client)
+        buffer = await client.download_media(target, in_memory=True)
+    return buffer.getvalue() if isinstance(buffer, BytesIO) else None