Harden TDLib large-media download path and address review follow-ups

docs/operator-workflows.md

@@ -145,7 +145,9 @@ The capability stays off until explicitly enabled. Rollout:
    ```
    Optional tuning: `TELEGRAM_TDLIB_SESSION_DIR` (default
    `~/.telegram-mcp-tdlib/main`), `TELEGRAM_TDLIB_DOWNLOAD_THRESHOLD_MB`
-   (default `20`).
+   (default `20`), `TELEGRAM_TDLIB_DB_ENCRYPTION_KEY` (default
+   `telegram-mcp-tdlib` — set a secret here so the TDLib database, which holds
+   the `main`-account auth key, isn't encrypted with a public constant).
 4. Watch telemetry (`download_post_backend` events: `backend`,
    `route_attempted`, `fallback_reason`) for the backend-used distribution
    and fallback rate before considering wider rollout (other accounts, lower
@@ -155,4 +157,7 @@ The capability stays off until explicitly enabled. Rollout:
 Every TDLib failure mode (session not authorized, network error, unsupported
 content type, lock not acquired within 5s) falls back to Telethon
 automatically using the connection already open for routing — there is no
-new failure mode a user can hit from this change.
+new failure mode a user can hit from this change. An unauthorized or revoked
+TDLib session is bounded by a readiness timeout (~15s) before it falls back,
+so a stale session adds a one-time delay per download, not a hang; re-run
+`scripts/tdlib_login.py` to clear it.

mcp/.env.example

@@ -22,6 +22,8 @@ TELEGRAM_API_HASH=
 # TELEGRAM_TDLIB_ENABLED=false
 # TELEGRAM_TDLIB_SESSION_DIR=~/.telegram-mcp-tdlib/main
 # TELEGRAM_TDLIB_DOWNLOAD_THRESHOLD_MB=20
+# Set a secret: the TDLib DB holds the account auth key; the default is a public constant.
+# TELEGRAM_TDLIB_DB_ENCRYPTION_KEY=change-me
 
 # Optional: HTTP transport settings for launchd/local daemon mode
 # TELEGRAM_MCP_TRANSPORT=streamable-http

mcp/src/telegram_mcp/download_post.py

@@ -123,6 +123,20 @@ def _telethon_media_size_bytes(msg) -> int | None:
     return None
 
 
+def _parse_threshold_mb(raw: str | None, *, default: float = 20.0) -> float:
+    """Parse the TDLib routing threshold from an env value, tolerating garbage.
+
+    A typo in TELEGRAM_TDLIB_DOWNLOAD_THRESHOLD_MB must not break downloads —
+    the Telethon path still works — so fall back to the default instead of
+    raising."""
+    if raw is None:
+        return default
+    try:
+        return float(raw)
+    except ValueError:
+        return default
+
+
 def _make_progress():
     last = {"pct": -5}
 
@@ -201,12 +215,13 @@ async def download_post(
         from . import tdlib_download  # lazy: keeps pytdbot fully optional at module load time
 
         tdlib_enabled = os.environ.get("TELEGRAM_TDLIB_ENABLED", "false").strip().lower() == "true"
-        threshold_mb = float(os.environ.get("TELEGRAM_TDLIB_DOWNLOAD_THRESHOLD_MB", "20"))
+        threshold_mb = _parse_threshold_mb(os.environ.get("TELEGRAM_TDLIB_DOWNLOAD_THRESHOLD_MB"))
+        media_size_bytes = _telethon_media_size_bytes(msg)
         route_to_tdlib = tdlib_download.should_route_to_tdlib(
             account=account,
             tdlib_enabled=tdlib_enabled,
             content_kind=_telethon_media_kind(msg),
-            media_size_bytes=_telethon_media_size_bytes(msg),
+            media_size_bytes=media_size_bytes,
             threshold_mb=threshold_mb,
         )
 
@@ -218,9 +233,16 @@ async def download_post(
             session_dir = Path(
                 os.environ.get("TELEGRAM_TDLIB_SESSION_DIR", "~/.telegram-mcp-tdlib/main")
             ).expanduser()
+            if not quiet:
+                print("PROGRESS routing large download through TDLib…", flush=True)
             try:
-                tdlib_path = await tdlib_download.download_via_tdlib(link=link, session_dir=session_dir)
-                shutil.copy2(tdlib_path, out)
+                await tdlib_download.download_via_tdlib(
+                    link=link,
+                    session_dir=session_dir,
+                    dest=out,
+                    expected_size_bytes=media_size_bytes,
+                    progress_callback=None if quiet else _make_progress(),
+                )
                 saved = str(out)
                 tdlib_backend_used = True
             except Exception as exc:  # noqa: BLE001 - any TDLib failure falls back to Telethon

mcp/src/telegram_mcp/locking.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import fcntl
 import os
 import time
@@ -53,14 +54,18 @@ def release(self) -> None:
             self._handle = None
 
 
-def try_acquire_with_timeout(
+async def try_acquire_with_timeout(
     lock: FileSessionLock,
     *,
     timeout_seconds: float = 5.0,
     poll_interval_seconds: float = 0.2,
 ) -> bool:
     """Retry ``lock.acquire()`` until it succeeds or ``timeout_seconds`` elapse.
 
+    Async so callers running inside an event loop (e.g. the TDLib download
+    path, which shares the loop with a live Telethon client) don't block it
+    while polling — the wait uses ``asyncio.sleep``.
+
     Returns True if acquired (caller owns the lock and must call release()).
     Returns False if the timeout elapsed without acquiring — the caller
     should fall back rather than treat this as an error.
@@ -73,4 +78,4 @@ def try_acquire_with_timeout(
         except SessionLockError:
             if time.monotonic() >= deadline:
                 return False
-            time.sleep(poll_interval_seconds)
+            await asyncio.sleep(poll_interval_seconds)

mcp/src/telegram_mcp/tdlib_download.py

@@ -9,7 +9,10 @@
 
 from __future__ import annotations
 
+import asyncio
 import os
+import shutil
+import time
 from pathlib import Path
 
 from .locking import FileSessionLock, try_acquire_with_timeout
@@ -18,6 +21,12 @@
 
 SUPPORTED_CONTENT_KINDS = frozenset({"video", "document", "photo", "audio"})
 
+DEFAULT_DB_ENCRYPTION_KEY = "telegram-mcp-tdlib"
+
+# Max wait for the TDLib session to reach authorizationStateReady after start().
+# An unauthorized/revoked session never gets there, so bound it and fall back.
+LOGIN_READY_TIMEOUT_SECONDS = 15.0
+
 
 class TdlibDownloadError(RuntimeError):
     """Raised for any TDLib download failure; callers should fall back to Telethon."""
@@ -36,9 +45,13 @@ def build_client(
     api_id: int,
     api_hash: str,
     files_directory: str,
-    database_encryption_key: str = "telegram-mcp-tdlib",
+    database_encryption_key: str | None = None,
 ):
     assert_isolated_from_telethon(files_directory)
+    if database_encryption_key is None:
+        database_encryption_key = os.environ.get(
+            "TELEGRAM_TDLIB_DB_ENCRYPTION_KEY", DEFAULT_DB_ENCRYPTION_KEY
+        )
     import pytdbot
 
     return pytdbot.Client(
@@ -72,9 +85,30 @@ def extract_file_id_from_message(message) -> int:
         return sizes[-1]["photo"]["id"]
     if content_type == "messageAudio":
         return content["audio"]["audio"]["id"]
+    if content_type == "messageAnimation":
+        return content["animation"]["animation"]["id"]
     raise ValueError(f"unsupported message content type for download: {content_type!r}")
 
 
+def assert_expected_size(*, downloaded_bytes: int, expected_bytes: int | None) -> None:
+    """Guard against TDLib resolving a different message than Telethon.
+
+    Telethon resolves the post by the chat/message_id parsed from the link; the
+    TDLib path re-resolves the same t.me URL via getMessageLinkInfo, which for a
+    link carrying a ?comment= / thread suffix can land on a *different* message
+    and silently download the wrong file. If the byte count we downloaded does
+    not match the size Telethon saw, refuse it so the caller falls back to the
+    correctly-resolved Telethon download. ``expected_bytes is None`` -> no check.
+    """
+    if expected_bytes is None:
+        return
+    if downloaded_bytes != expected_bytes:
+        raise TdlibDownloadError(
+            f"TDLib downloaded {downloaded_bytes} bytes but Telethon expected "
+            f"{expected_bytes}; the link likely resolved to a different message"
+        )
+
+
 def should_route_to_tdlib(
     *,
     account: str,
@@ -94,12 +128,84 @@ def should_route_to_tdlib(
     return media_size_bytes >= threshold_mb * 1024 * 1024
 
 
-async def download_via_tdlib(*, link: str, session_dir: Path) -> Path:
-    """Resolve `link` via TDLib and download it fully. Returns the local file
-    path on success. Raises TdlibDownloadError on any failure (lock timeout,
-    TDLib error, incomplete download) — the caller decides to fall back."""
+async def start_client_ready(
+    client,
+    *,
+    timeout_seconds: float = LOGIN_READY_TIMEOUT_SECONDS,
+    poll_interval_seconds: float = 0.25,
+) -> None:
+    """Start ``client`` and wait until it reaches authorizationStateReady.
+
+    ``pytdbot.Client.start()`` only auto-drives login for bot-token clients; for
+    a user-account client (as built here) the default ``wait_login=True`` blocks
+    forever when the session is unauthorized, because the event it waits on is
+    only set on authorizationStateReady. So start with ``wait_login=False`` and
+    poll the state ourselves, raising ``TdlibDownloadError`` on timeout so the
+    caller falls back to Telethon instead of hanging.
+    """
+    await client.start(wait_login=False)
+    deadline = time.monotonic() + timeout_seconds
+    while True:
+        state = client.authorization_state
+        if state == "authorizationStateReady":
+            return
+        if time.monotonic() >= deadline:
+            raise TdlibDownloadError(
+                f"TDLib session not ready (state={state!r}); "
+                "run scripts/tdlib_login.py to (re)authorize"
+            )
+        await asyncio.sleep(poll_interval_seconds)
+
+
+async def _download_with_progress(
+    client,
+    *,
+    file_id: int,
+    progress_callback,
+    total_bytes: int | None,
+    poll_interval_seconds: float = 1.0,
+):
+    """Start an async download and poll getFile until it completes, reporting
+    progress each poll. TDLib's ``synchronous=True`` download stays silent until
+    it finishes, which for a multi-gig file means minutes with no feedback; the
+    poll loop lets the caller surface real progress. Returns the completed File.
+    """
+    raise_if_error(
+        await client.downloadFile(
+            file_id=file_id, priority=1, synchronous=False, offset=0, limit=0
+        )
+    )
+    while True:
+        file = raise_if_error(await client.getFile(file_id=file_id))
+        local = file["local"]
+        downloaded = int(local["downloaded_size"]) if local else 0
+        total = total_bytes or int(file["size"]) or downloaded
+        progress_callback(downloaded, total)
+        if local and bool(local["is_downloading_completed"]):
+            return file
+        await asyncio.sleep(poll_interval_seconds)
+
+
+async def download_via_tdlib(
+    *,
+    link: str,
+    session_dir: Path,
+    dest: Path,
+    expected_size_bytes: int | None = None,
+    progress_callback=None,
+) -> Path:
+    """Resolve `link` via TDLib, download it fully, copy it to `dest`, and drop
+    the TDLib-side copy. Returns `dest` on success. Raises TdlibDownloadError on
+    any failure (lock timeout, unauthorized session, TDLib error, incomplete
+    download, size mismatch) — the caller decides to fall back.
+
+    ``expected_size_bytes`` is the size Telethon saw for the resolved message;
+    it guards against the link resolving to a different message on the TDLib
+    side (see ``assert_expected_size``) and doubles as the progress denominator.
+    When ``progress_callback`` is given the download runs asynchronously with a
+    progress poll; otherwise it blocks synchronously."""
     lock = FileSessionLock(session_dir / "download.lock")
-    if not try_acquire_with_timeout(lock, timeout_seconds=5.0):
+    if not await try_acquire_with_timeout(lock, timeout_seconds=5.0):
         raise TdlibDownloadError("could not acquire TDLib session lock within 5s")
 
     try:
@@ -108,7 +214,7 @@ async def download_via_tdlib(*, link: str, session_dir: Path) -> Path:
             api_hash=os.environ["TELEGRAM_API_HASH"],
             files_directory=str(session_dir),
         )
-        await client.start()
+        await start_client_ready(client)
         try:
             link_info = raise_if_error(await client.getMessageLinkInfo(url=link))
             message = raise_if_error(
@@ -117,15 +223,35 @@ async def download_via_tdlib(*, link: str, session_dir: Path) -> Path:
                 )
             )
             file_id = extract_file_id_from_message(message)
-            result = raise_if_error(
-                await client.downloadFile(
-                    file_id=file_id, priority=1, synchronous=True, offset=0, limit=0
+            if progress_callback is None:
+                result = raise_if_error(
+                    await client.downloadFile(
+                        file_id=file_id, priority=1, synchronous=True, offset=0, limit=0
+                    )
+                )
+            else:
+                result = await _download_with_progress(
+                    client,
+                    file_id=file_id,
+                    progress_callback=progress_callback,
+                    total_bytes=expected_size_bytes,
                 )
-            )
             local = result["local"]
             if not local or not bool(local["is_downloading_completed"]):
                 raise TdlibDownloadError("TDLib download did not complete")
-            return Path(local["path"])
+            assert_expected_size(
+                downloaded_bytes=int(local["downloaded_size"]),
+                expected_bytes=expected_size_bytes,
+            )
+            shutil.copy2(local["path"], dest)
+            # Drop the TDLib-side copy so files_directory doesn't grow without
+            # bound; resumability matters for partial files, not ones we've
+            # already copied out. Best-effort — a cleanup failure isn't fatal.
+            try:
+                raise_if_error(await client.deleteFile(file_id=file_id))
+            except Exception:  # noqa: BLE001 - cleanup is best-effort
+                pass
+            return dest
         finally:
             await client.stop()
     finally:

mcp/tests/test_download_post.py

@@ -4,6 +4,7 @@
 from telegram_mcp.download_post import (
     ParsedLink,
     _ext_from_message,
+    _parse_threshold_mb,
     _telethon_media_kind,
     _telethon_media_size_bytes,
     parse_post_link,
@@ -118,5 +119,19 @@ def test_no_media_is_none(self):
         self.assertIsNone(_telethon_media_size_bytes(msg))
 
 
+class ParseThresholdMbTests(unittest.TestCase):
+    def test_valid_value(self):
+        self.assertEqual(_parse_threshold_mb("50"), 50.0)
+
+    def test_none_falls_back_to_default(self):
+        self.assertEqual(_parse_threshold_mb(None), 20.0)
+
+    def test_garbage_falls_back_to_default_instead_of_raising(self):
+        self.assertEqual(_parse_threshold_mb("twenty"), 20.0)
+
+    def test_empty_falls_back_to_default(self):
+        self.assertEqual(_parse_threshold_mb(""), 20.0)
+
+
 if __name__ == "__main__":
     unittest.main()

mcp/tests/test_locking.py

@@ -20,27 +20,29 @@ def test_file_session_lock_rejects_second_holder(self):
             finally:
                 first.release()
 
-    def test_try_acquire_with_timeout_returns_false_when_held(self):
+
+class TryAcquireWithTimeoutTests(unittest.IsolatedAsyncioTestCase):
+    async def test_returns_false_when_held(self):
         with tempfile.TemporaryDirectory() as tmp:
             lock_path = Path(tmp) / "telegram.lock"
 
             holder = FileSessionLock(lock_path)
             holder.acquire()
             try:
                 waiter = FileSessionLock(lock_path)
-                acquired = try_acquire_with_timeout(
+                acquired = await try_acquire_with_timeout(
                     waiter, timeout_seconds=0.3, poll_interval_seconds=0.1
                 )
                 self.assertFalse(acquired)
             finally:
                 holder.release()
 
-    def test_try_acquire_with_timeout_succeeds_once_free(self):
+    async def test_succeeds_once_free(self):
         with tempfile.TemporaryDirectory() as tmp:
             lock_path = Path(tmp) / "telegram.lock"
 
             lock = FileSessionLock(lock_path)
-            acquired = try_acquire_with_timeout(lock, timeout_seconds=0.3)
+            acquired = await try_acquire_with_timeout(lock, timeout_seconds=0.3)
             try:
                 self.assertTrue(acquired)
             finally: