From 9c12628304f8d0813a667dea5d8c0ef15fe9a768 Mon Sep 17 00:00:00 2001 From: h Date: Tue, 1 Sep 2026 21:05:48 +0200 Subject: [PATCH] fix(userbot): drop stale media sessions and survive download timeouts --- .../src/userbot/modules/avatars/downloader.py | 13 +++-- backend/src/userbot/modules/download.py | 56 +++++++++++++++++++ .../modules/jobs/handlers/fetch_avatar.py | 8 +-- .../jobs/handlers/fetch_custom_emoji.py | 8 +-- .../src/userbot/modules/media/downloader.py | 14 +++-- .../src/userbot/modules/stories/service.py | 8 +-- 6 files changed, 83 insertions(+), 24 deletions(-) create mode 100644 backend/src/userbot/modules/download.py diff --git a/backend/src/userbot/modules/avatars/downloader.py b/backend/src/userbot/modules/avatars/downloader.py index 21df1b5..7e08ee5 100644 --- a/backend/src/userbot/modules/avatars/downloader.py +++ b/backend/src/userbot/modules/avatars/downloader.py @@ -1,9 +1,9 @@ -from io import BytesIO - from pyrogram import Client from userbot.modules.avatars import repository from userbot.modules.capture.context import CaptureContext +from userbot.modules.download import download_bytes +from utils.logging import logger async def capture_avatar( # noqa: PLR0913 @@ -20,9 +20,12 @@ async def capture_avatar( # noqa: PLR0913 storage_key: str | None = None file_size: int | None = None downloaded = False - buffer = await client.download_media(file_id, in_memory=True) - if isinstance(buffer, BytesIO): - data = buffer.getvalue() + try: + data = await download_bytes(client, file_id) + except TimeoutError: + logger.warning(f"[yellow]Avatar download timed out for {owner_id}.[/]") + data = None + if data is not None: storage_key = ctx.storage.put(data) file_size = len(data) downloaded = True diff --git a/backend/src/userbot/modules/download.py b/backend/src/userbot/modules/download.py new file mode 100644 index 0000000..c4cfc46 --- /dev/null +++ b/backend/src/userbot/modules/download.py @@ -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 diff --git a/backend/src/userbot/modules/jobs/handlers/fetch_avatar.py b/backend/src/userbot/modules/jobs/handlers/fetch_avatar.py index bc7ed0d..a95c2ed 100644 --- a/backend/src/userbot/modules/jobs/handlers/fetch_avatar.py +++ b/backend/src/userbot/modules/jobs/handlers/fetch_avatar.py @@ -1,10 +1,9 @@ -from io import BytesIO - from pyrogram import Client from pyrogram.errors import FileIdInvalid, FileReferenceExpired, FileReferenceInvalid from pyrogram.types import Photo from userbot.modules.avatars.repository import get_avatar_file, mark_avatar_downloaded +from userbot.modules.download import download_bytes from userbot.modules.jobs.context import JobContext from userbot.modules.jobs.registry import register @@ -29,13 +28,12 @@ async def _download( client: Client, owner_id: int, unique_id: str, file_id: str ) -> bytes | None: try: - buffer = await client.download_media(file_id, in_memory=True) + return await download_bytes(client, file_id) except STALE_FILE_ID: fresh = await _fresh_file_id(client, owner_id, unique_id) if fresh is None: return None - buffer = await client.download_media(fresh, in_memory=True) - return buffer.getvalue() if isinstance(buffer, BytesIO) else None + return await download_bytes(client, fresh) @register("fetch_avatar") diff --git a/backend/src/userbot/modules/jobs/handlers/fetch_custom_emoji.py b/backend/src/userbot/modules/jobs/handlers/fetch_custom_emoji.py index 96c903c..ad0e1e9 100644 --- a/backend/src/userbot/modules/jobs/handlers/fetch_custom_emoji.py +++ b/backend/src/userbot/modules/jobs/handlers/fetch_custom_emoji.py @@ -1,8 +1,7 @@ -from io import BytesIO - from pyrogram.types import Sticker from userbot.modules.custom_emoji.repository import is_downloaded, upsert_downloaded +from userbot.modules.download import download_bytes from userbot.modules.jobs.context import JobContext from userbot.modules.jobs.registry import register @@ -30,10 +29,9 @@ async def fetch_custom_emoji(ctx: JobContext) -> None: if not stickers: return sticker = stickers[0] - buffer = await client.download_media(sticker.file_id, in_memory=True) - if not isinstance(buffer, BytesIO): + data = await download_bytes(client, sticker.file_id) + if data is None: return - data = buffer.getvalue() storage_key = capture.storage.put(data) await upsert_downloaded( ctx.pool, diff --git a/backend/src/userbot/modules/media/downloader.py b/backend/src/userbot/modules/media/downloader.py index fa70418..eacc32c 100644 --- a/backend/src/userbot/modules/media/downloader.py +++ b/backend/src/userbot/modules/media/downloader.py @@ -1,4 +1,3 @@ -from io import BytesIO from typing import Any from pyrogram import Client @@ -6,6 +5,8 @@ from pyrogram.types import Message from userbot.modules.capture import repository from userbot.modules.capture.context import CaptureContext +from userbot.modules.download import download_bytes +from utils.logging import logger from utils.policy.models import CaptureToggles _MEDIA_ATTRS = ( @@ -80,9 +81,14 @@ async def capture_media( # noqa: PLR0913 downloaded = True else: target = message if getattr(message, kind or "", None) is obj else obj - buffer = await client.download_media(target, in_memory=True) - if isinstance(buffer, BytesIO): - data = buffer.getvalue() + try: + data = await download_bytes(client, target) + except TimeoutError: + logger.warning( + f"[yellow]Media download timed out for {chat_id}/{message_id}.[/]" + ) + data = None + if data is not None: storage_key = ctx.storage.put(data) file_size = len(data) downloaded = True diff --git a/backend/src/userbot/modules/stories/service.py b/backend/src/userbot/modules/stories/service.py index 5dff4c0..52f4fca 100644 --- a/backend/src/userbot/modules/stories/service.py +++ b/backend/src/userbot/modules/stories/service.py @@ -1,9 +1,8 @@ -from io import BytesIO - from pyrogram import Client from pyrogram.types import Story from userbot.modules.capture.context import CaptureContext +from userbot.modules.download import download_bytes from userbot.modules.stories import repository @@ -24,9 +23,8 @@ async def save_story(client: Client, capture: CaptureContext, story: Story) -> N 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() + data = await download_bytes(client, story) + if data is not None: storage_key = capture.storage.put(data) file_size = len(data) downloaded = True