fix(stability): stage 1 — timeouts on all TG RPC + resilient download worker + supervision

Eliminates the app-wide hangs where one stuck Telegram RPC holds the single
tg_rpc gate permit forever (freezing every RSS/HTML request) and where the
background download worker dies silently.

1.1 Every RPC under the global gate is now bounded by asyncio.wait_for
    (config tg_rpc_timeout, env TG_RPC_TIMEOUT, default 60). The wait_for wraps
    ONLY the RPC body, never the gate acquire, so queue backpressure stays
    legitimate; the paginated get_chat_history is collected in an inner coroutine
    so the async-for can be bounded. A timeout propagates out of the gate so its
    permit is released (no leak).
1.2 The other live RPC paths get the same treatment: _reply_enrichment's
    get_messages now runs under the gate + timeout; PostParser.get_post is bounded
    at 30s (and a stray print removed); /health's get_me at 10s.
1.3 background_download_worker: queue.get() moved out of the try so task_done()
    in finally balances exactly one get (no ValueError that killed the worker);
    FloodWait is caught BEFORE the generic Exception (sleeps, does not flood);
    download_new_files uses put_nowait so a full queue no longer blocks the
    sweeper forever.
1.4 Both background tasks run under a _supervised() wrapper: a crash or an
    unexpected return is logged CRITICAL and restarted (restarts rate-limited to
    once per 60s so a hard-failing task can't spin), while CancelledError is
    propagated to the child for a clean shutdown.

Tests (tests/test_stage1_hangs.py): gate timeout releases the permit + a second
call succeeds; gate cancel mid-spacing releases the permit; the worker survives
Exception and FloodWait with balanced task_done; and _supervised restarts a
crashing task (rate-limited) and an unexpected return, and propagates cancellation
to its child. 180 passed (174 baseline + 6).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

api_server.py

@@ -72,6 +72,37 @@ HTTP_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(3)  # semaphore for live HTTP media
 BACKGROUND_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(2)  # semaphore for background cache worker
 download_queue = asyncio.Queue(maxsize=100)
 
+async def _supervised(factory, name: str, min_restart_interval: float = 60.0):
+    """Run factory() forever, restarting it if it dies with a non-cancellation error.
+
+    A background loop that returns or raises (anything except CancelledError) is logged
+    at CRITICAL and restarted, but successive (re)starts are spaced at least
+    min_restart_interval seconds apart so a hard-failing task can't spin the event loop.
+    CancelledError (shutdown) is propagated to the child and stops supervision.
+    """
+    while True:
+        start = time.monotonic()
+        task = asyncio.create_task(factory(), name=name)
+        try:
+            await task
+            # A supervised background loop is not expected to return on its own.
+            logger.critical(f"supervised_task_exited: {name} returned unexpectedly; restarting")
+        except asyncio.CancelledError:
+            # Shutdown or external cancel: propagate to the child, then stop supervising.
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+            raise
+        except Exception as e:
+            logger.critical(f"supervised_task_crashed: {name} died with {e!r}; restarting", exc_info=True)
+        # Rate-limit restarts: keep successive starts at least min_restart_interval apart.
+        elapsed = time.monotonic() - start
+        if elapsed < min_restart_interval:
+            await asyncio.sleep(min_restart_interval - elapsed)
+
+
 @asynccontextmanager
 async def lifespan(_: FastAPI):
     setup_logging(Config["log_level"])
@@ -84,8 +115,10 @@ async def lifespan(_: FastAPI):
     await asyncio.to_thread(init_db_sync, DB_PATH)
 
     await client.start()
-    background_task = asyncio.create_task(cache_media_files()) # Start background task
-    worker_task = asyncio.create_task(background_download_worker()) # Start download worker
+    # Supervise the background tasks: if either dies (not via cancellation) it is logged
+    # CRITICAL and restarted, so a crash can no longer silently stop cache sweeping or downloads.
+    background_task = asyncio.create_task(_supervised(cache_media_files, "cache_media_files"))
+    worker_task = asyncio.create_task(_supervised(background_download_worker, "background_download_worker"))
     yield
     background_task.cancel() # Cleanup
     worker_task.cancel()
@@ -602,7 +635,10 @@ async def download_new_files(media_files: list, cache_dir: str) -> None:
             cache_path = os.path.join(post_dir, file_unique_id)
             if not os.path.exists(cache_path):
                 try:
-                    await download_queue.put((channel, post_id, file_unique_id))
+                    # put_nowait so a full queue raises QueueFull instead of blocking
+                    # cache_media_files (and thus the sweeper) forever. `await put()`
+                    # never raises QueueFull, which made the except below dead code.
+                    download_queue.put_nowait((channel, post_id, file_unique_id))
                     files_queued += 1
                     logger.debug(f"Queued for background download: {channel}/{post_id}/{file_unique_id}")
                 except asyncio.QueueFull:
@@ -620,19 +656,22 @@ async def download_new_files(media_files: list, cache_dir: str) -> None:
 async def background_download_worker():
     """Worker that processes downloads from queue"""
     while True:
+        # get() is OUTSIDE the try so task_done() in finally always balances exactly one
+        # successful get(). Cancellation propagates cleanly here (nothing to unbalance).
+        item = await download_queue.get()
+        channel, post_id, file_unique_id = item
+        logger.info(f"Background download: {channel}/{post_id}/{file_unique_id}")
         try:
-            channel, post_id, file_unique_id = await download_queue.get()
-            logger.info(f"Background download: {channel}/{post_id}/{file_unique_id}")
-            
-            try:
-                async with BACKGROUND_DOWNLOAD_SEMAPHORE:  # limit concurrent background downloads
-                    await download_media_file(channel, post_id, file_unique_id)
-                await asyncio.sleep(2)
-            except Exception as e:
-                logger.error(f"Background download error for {channel}/{post_id}/{file_unique_id}: {e}")
-                
+            async with BACKGROUND_DOWNLOAD_SEMAPHORE:  # limit concurrent background downloads
+                await download_media_file(channel, post_id, file_unique_id)
+            await asyncio.sleep(2)
+        except errors.FloodWait as e:
+            # Must be caught BEFORE the generic Exception (FloodWait subclasses RPCError),
+            # otherwise the worker would hammer Telegram while under a flood wait.
+            logger.warning(f"bg_download_floodwait: {channel}/{post_id}/{file_unique_id} sleeping {e.value}s")
+            await asyncio.sleep(min(int(e.value) + 5, 900))
         except Exception as e:
-            logger.error(f"Background download worker error: {e}")
+            logger.error(f"Background download error for {channel}/{post_id}/{file_unique_id}: {e}")
         finally:
             download_queue.task_done()
 
@@ -903,8 +942,9 @@ async def health_check(request: Request, token: str | None = None) -> Response:
         logger.info(f"Local request, skipping token check for health check.")
             
     try:
-        me = await client.client.get_me()
-        
+        # Bound the Telegram RPC so a hung get_me cannot hang the healthcheck.
+        me = await asyncio.wait_for(client.client.get_me(), timeout=10)
+
         # Offload heavy filesystem scanning to threadpool
         cache_stats = await asyncio.to_thread(calculate_cache_stats)
         

config.py

@@ -112,6 +112,9 @@ def get_settings() -> dict[str, Any]:
         "show_bridge_link": os.getenv("SHOW_BRIDGE_LINK", "False").strip() in ["True", "true"],
         "show_post_flags": os.getenv("SHOW_POST_FLAGS", "False").strip() in ["True", "true"],
         "proxy": proxy,
+        # Hard cap (seconds) on any single live Telegram RPC held under the global RPC gate,
+        # so a hung MTProto call can never pin the gate (and the whole app) indefinitely.
+        "tg_rpc_timeout": _parse_int_env("TG_RPC_TIMEOUT", 60),
         "tg_watchdog_enabled": os.getenv("TG_WATCHDOG_ENABLED", "true").strip().lower() not in ["false", "0", "no", "off", "disable", "disabled"],
         "tg_watchdog_interval": _parse_int_env("TG_WATCHDOG_INTERVAL", 60),
         "tg_watchdog_timeout": _parse_int_env("TG_WATCHDOG_TIMEOUT", 10),

post_parser.py

@@ -92,10 +92,13 @@ class PostParser:
                         post_id: int, 
                         output_type: str = 'json', 
                         debug: bool = False) -> Union[str, Dict[Any, Any], None]:
-        print(f"Getting post {channel}, {post_id}")
         try:
             prepared_channel_id: Union[str, int] = self.channel_name_prepare(channel)
-            message = await self.client.get_messages(prepared_channel_id, post_id)
+            # Bound the single-post fetch so a hung RPC cannot block the request forever.
+            message = await asyncio.wait_for(
+                self.client.get_messages(prepared_channel_id, post_id),
+                timeout=30,
+            )
 
             if Config["debug"]: print(message)
 

rss_generator.py

@@ -22,6 +22,7 @@ from pyrogram import errors, Client
 from pyrogram.types import Message
 from post_parser import PostParser
 from config import get_settings
+from tg_throttle import tg_rpc
 from bleach.css_sanitizer import CSSSanitizer
 from bleach import clean as HTMLSanitizer
 
@@ -502,7 +503,13 @@ async def _reply_enrichment(client: Client, messages: list[Message]) -> list[Mes
     for chat_id, chat_msgs in chat_messages.items():
         ids_to_fetch = [m.id for m in chat_msgs]
         try:
-            fetched = await client.get_messages(chat_id, ids_to_fetch)
+            # Throttle under the global RPC gate and bound the call so a hung
+            # get_messages cannot pin the gate (wait_for wraps the RPC, not the gate entry).
+            async with tg_rpc():
+                fetched = await asyncio.wait_for(
+                    client.get_messages(chat_id, ids_to_fetch),
+                    timeout=Config["tg_rpc_timeout"],
+                )
             # get_messages may return a single Message or a list
             if not isinstance(fetched, list):
                 fetched = [fetched]

tests/mock_config.py

@@ -1,6 +1,10 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
+def setup_logging(level_name: str = "INFO") -> None:
+    """No-op logging setup for tests (mirrors config.setup_logging signature)."""
+    return None
+
 def get_settings():
     """
     Mock config for testing without requiring TG_API_ID and TG_API_HASH
@@ -20,6 +24,7 @@ def get_settings():
         "show_post_flags": True,
         "proxy": None,
         "trusted_proxies": [],
+        "tg_rpc_timeout": 60,
         "tg_watchdog_enabled": True,
         "tg_watchdog_interval": 60,
         "tg_watchdog_timeout": 10,

tests/test_stage1_hangs.py

@@ -0,0 +1,214 @@
+# flake8: noqa
+# pylint: disable=protected-access, missing-function-docstring, missing-class-docstring
+# pylint: disable=redefined-outer-name, logging-fstring-interpolation, line-too-long
+# pylance: disable=reportMissingImports, reportMissingModuleSource
+"""
+Stage 1 (anti-hang) regression tests:
+- RPC gate timeout: a hung RPC times out and the gate permit is NOT leaked; a
+  subsequent call still succeeds.
+- Background worker: a download that raises Exception / FloodWait does not kill the
+  worker, and task_done stays balanced so queue.join() completes.
+- Gate cancellation during the spacing wait does not lose the permit.
+"""
+import os
+import sys
+import time
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+# Add project root to sys.path and mock the config module (same pattern as the other tests).
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+sys.modules['config'] = __import__('tests.mock_config', fromlist=['get_settings'])
+
+import tg_throttle
+import tg_cache
+from pyrogram import errors
+
+
+# --------------------------------------------------------------------------- #
+# 1.1 — RPC gate timeout releases the permit; second call succeeds.
+# --------------------------------------------------------------------------- #
+@pytest.mark.asyncio
+async def test_gate_timeout_releases_permit(monkeypatch):
+    # Short timeout so the hung RPC fails fast, and no min-interval spacing to slow the test.
+    monkeypatch.setitem(tg_cache.Config, "tg_rpc_timeout", 0.1)
+    monkeypatch.setattr(tg_throttle, "_MIN_INTERVAL", 0.0)
+    # Bypass the on-disk chat cache so we always hit the live RPC path.
+    monkeypatch.setattr(tg_cache, "_get_chat_from_cache", lambda *a, **k: None)
+    monkeypatch.setattr(tg_cache, "_save_chat_to_cache", lambda *a, **k: None)
+
+    permits_before = tg_throttle._sem._value
+    never = asyncio.Event()  # never set -> RPC hangs forever
+
+    class HungClient:
+        async def get_chat(self, channel_id):
+            await never.wait()
+
+    # First call: the RPC hangs and must time out (not hang the whole app).
+    with pytest.raises(asyncio.TimeoutError):
+        await tg_cache.cached_get_chat(HungClient(), "hung_channel")
+
+    # The permit was released on timeout (gate fully available again) — no leak.
+    assert tg_throttle._sem._value == permits_before
+
+    # A second call still goes through the gate and succeeds.
+    class OkClient:
+        async def get_chat(self, channel_id):
+            return SimpleNamespace(id=42, title="ok", username="okchan")
+
+    res = await tg_cache.cached_get_chat(OkClient(), "ok_channel")
+    assert res.id == 42
+    # Permit released again after the successful call.
+    assert tg_throttle._sem._value == permits_before
+
+
+# --------------------------------------------------------------------------- #
+# 1.4 — Gate cancellation during the spacing wait does not lose the permit.
+# --------------------------------------------------------------------------- #
+@pytest.mark.asyncio
+async def test_gate_cancel_during_spacing_releases_permit(monkeypatch):
+    # Force a spacing wait long enough to cancel inside it.
+    monkeypatch.setattr(tg_throttle, "_MIN_INTERVAL", 1.0)
+    monkeypatch.setattr(tg_throttle, "_last_start", time.monotonic())
+
+    permits_before = tg_throttle._sem._value
+
+    async def enter_gate():
+        async with tg_throttle.tg_rpc():
+            pass
+
+    task = asyncio.create_task(enter_gate())
+    # Let it acquire the semaphore and start sleeping for the spacing interval.
+    await asyncio.sleep(0.05)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    # Cancelled mid-spacing: the acquired permit must be returned.
+    assert tg_throttle._sem._value == permits_before
+
+
+# --------------------------------------------------------------------------- #
+# 1.3 — Background worker survives download errors; task_done stays balanced.
+# --------------------------------------------------------------------------- #
+@pytest.mark.asyncio
+async def test_worker_survives_errors_and_balances_task_done(monkeypatch):
+    import api_server
+
+    # Speed up the worker's post-download / flood-wait sleeps.
+    async def _fast_sleep(*_a, **_k):
+        return
+    monkeypatch.setattr(api_server.asyncio, "sleep", _fast_sleep)
+
+    # Fresh queue so we don't interfere with (or depend on) module state.
+    q = asyncio.Queue(maxsize=100)
+    monkeypatch.setattr(api_server, "download_queue", q)
+
+    processed = []
+
+    async def fake_download(channel, post_id, file_unique_id):
+        processed.append(file_unique_id)
+        if file_unique_id == "boom":
+            raise RuntimeError("simulated download failure")
+        if file_unique_id == "flood":
+            raise errors.FloodWait(value=1)
+        # "ok" succeeds
+
+    monkeypatch.setattr(api_server, "download_media_file", fake_download)
+
+    worker = asyncio.create_task(api_server.background_download_worker())
+    try:
+        for fid in ("boom", "flood", "ok"):
+            q.put_nowait(("chan", 1, fid))
+
+        # If the worker died on the first error, or task_done() were unbalanced,
+        # join() would never complete and this wait_for would time out.
+        await asyncio.wait_for(q.join(), timeout=5)
+
+        # The worker stayed alive across the Exception and the FloodWait and drained all items.
+        assert processed == ["boom", "flood", "ok"]
+    finally:
+        worker.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await worker
+
+
+# --------------------------------------------------------------------------- #
+# 1.4 — _supervised: restart on crash (rate-limited), restart on unexpected
+# return, and clean cancellation propagation on shutdown.
+# --------------------------------------------------------------------------- #
+@pytest.mark.asyncio
+async def test_supervised_restarts_crashing_task_rate_limited():
+    import api_server
+
+    starts = []
+
+    async def crashing():
+        starts.append(time.monotonic())
+        raise RuntimeError("always dies")
+
+    # Tiny min interval so the test is fast, but non-zero so we can assert spacing.
+    sup = asyncio.create_task(
+        api_server._supervised(crashing, "crasher", min_restart_interval=0.05)
+    )
+    # Let it crash-and-restart a few times.
+    await asyncio.sleep(0.28)
+    sup.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await sup
+
+    # It restarted several times (did not give up after the first crash)...
+    assert len(starts) >= 3
+    # ...but restarts were spaced at least ~min_restart_interval apart (no spin).
+    gaps = [b - a for a, b in zip(starts, starts[1:])]
+    assert all(g >= 0.045 for g in gaps), gaps
+
+
+@pytest.mark.asyncio
+async def test_supervised_restarts_on_unexpected_return():
+    import api_server
+
+    runs = []
+
+    async def returns_immediately():
+        runs.append(1)
+        # A supervised background loop is not meant to return; _supervised must
+        # log CRITICAL and restart it rather than stop supervising.
+        return
+
+    sup = asyncio.create_task(
+        api_server._supervised(returns_immediately, "returner", min_restart_interval=0.02)
+    )
+    await asyncio.sleep(0.1)
+    sup.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await sup
+
+    assert len(runs) >= 2  # restarted after the unexpected return
+
+
+@pytest.mark.asyncio
+async def test_supervised_cancellation_propagates_to_child():
+    import api_server
+
+    child_cancelled = asyncio.Event()
+
+    async def long_running():
+        try:
+            await asyncio.Event().wait()  # runs until cancelled
+        except asyncio.CancelledError:
+            child_cancelled.set()
+            raise
+
+    sup = asyncio.create_task(
+        api_server._supervised(long_running, "longrun")
+    )
+    await asyncio.sleep(0.05)  # let the child start
+    sup.cancel()
+    # Cancelling the supervisor must propagate CancelledError out (clean shutdown)...
+    with pytest.raises(asyncio.CancelledError):
+        await sup
+    # ...and the child task must have been cancelled too (no leaked background task).
+    assert child_cancelled.is_set()

tg_cache.py

@@ -22,9 +22,12 @@ from typing import Any, Optional, Union, List
 from pyrogram import Client
 from pyrogram.types import Chat, Message
 from tg_throttle import tg_rpc
+from config import get_settings
 
 logger = logging.getLogger(__name__)
 
+Config = get_settings()
+
 # Path to cache directory
 CACHE_DIR = os.path.join('data', 'tgcache')
 
@@ -138,10 +141,16 @@ async def cached_get_chat_history(client: Client, channel_id: Union[str, int], l
     
     try:
         logger.info(f"history_cache_request: fetching fresh history for channel {channel_id}, limit {limit}")
-        messages = []
+        # Hold the global RPC gate only for the live fetch. The paginated `async for`
+        # cannot be wrapped in wait_for directly, so collect it in an inner coroutine
+        # and bound THAT with the timeout. wait_for wraps only the RPC body, never the
+        # gate entry (`async with tg_rpc()`), so legitimate queue backpressure is not
+        # mistaken for a hang.
         async with tg_rpc():
-            async for message in client.get_chat_history(channel_id, limit=limit):
-                messages.append(message)
+            async def _collect():
+                # Full paginated history; bounded by the outer wait_for.
+                return [m async for m in client.get_chat_history(channel_id, limit=limit)]
+            messages = await asyncio.wait_for(_collect(), timeout=Config["tg_rpc_timeout"])
         await asyncio.to_thread(_save_history_to_cache, channel_id, messages, limit)
         
         return messages
@@ -214,7 +223,8 @@ async def cached_get_chat(client: Client, channel_id: Union[str, int]) -> Simple
 
     logger.info(f"chatinfo_cache_request: fetching fresh chat info for channel {channel_id}")
     async with tg_rpc():
-        chat = await client.get_chat(channel_id)
+        # Bound the RPC itself, not the gate entry above it.
+        chat = await asyncio.wait_for(client.get_chat(channel_id), timeout=Config["tg_rpc_timeout"])
 
     data = {
         'id': getattr(chat, 'id', None),