diff --git a/docs/operator-workflows.md b/docs/operator-workflows.md index 038c44c..e6443fb 100644 --- a/docs/operator-workflows.md +++ b/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. diff --git a/mcp/.env.example b/mcp/.env.example index 834426c..d68d347 100644 --- a/mcp/.env.example +++ b/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 diff --git a/mcp/src/telegram_mcp/download_post.py b/mcp/src/telegram_mcp/download_post.py index e41bfc7..323ace0 100644 --- a/mcp/src/telegram_mcp/download_post.py +++ b/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 diff --git a/mcp/src/telegram_mcp/locking.py b/mcp/src/telegram_mcp/locking.py index 7312b32..30fc7bb 100644 --- a/mcp/src/telegram_mcp/locking.py +++ b/mcp/src/telegram_mcp/locking.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import fcntl import os import time @@ -53,7 +54,7 @@ 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, @@ -61,6 +62,10 @@ def try_acquire_with_timeout( ) -> 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) diff --git a/mcp/src/telegram_mcp/tdlib_download.py b/mcp/src/telegram_mcp/tdlib_download.py index 2dc55c2..5a783ab 100644 --- a/mcp/src/telegram_mcp/tdlib_download.py +++ b/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: diff --git a/mcp/tests/test_download_post.py b/mcp/tests/test_download_post.py index 937a2f0..6d5130f 100644 --- a/mcp/tests/test_download_post.py +++ b/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() diff --git a/mcp/tests/test_locking.py b/mcp/tests/test_locking.py index 2261ee9..8f56e8e 100644 --- a/mcp/tests/test_locking.py +++ b/mcp/tests/test_locking.py @@ -20,7 +20,9 @@ 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" @@ -28,19 +30,19 @@ def test_try_acquire_with_timeout_returns_false_when_held(self): 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: diff --git a/mcp/tests/test_tdlib_download.py b/mcp/tests/test_tdlib_download.py index bad309b..6732817 100644 --- a/mcp/tests/test_tdlib_download.py +++ b/mcp/tests/test_tdlib_download.py @@ -1,11 +1,41 @@ import importlib.util import unittest -from telegram_mcp.tdlib_download import assert_isolated_from_telethon, should_route_to_tdlib +from telegram_mcp.tdlib_download import ( + TdlibDownloadError, + assert_expected_size, + assert_isolated_from_telethon, + should_route_to_tdlib, + start_client_ready, +) PYTDBOT_AVAILABLE = importlib.util.find_spec("pytdbot") is not None +class _FakeClient: + """Minimal stand-in for pytdbot.Client for start_client_ready tests. + + start_client_ready never imports pytdbot itself — it only calls + ``start(wait_login=False)`` and reads ``authorization_state`` — so it can + be exercised without the optional extra installed. + """ + + def __init__(self, states): + # states: list of authorization_state values yielded on each read; + # the last value repeats once exhausted. + self._states = list(states) + self.start_calls = [] + + async def start(self, wait_login=True): + self.start_calls.append(wait_login) + + @property + def authorization_state(self): + if len(self._states) > 1: + return self._states.pop(0) + return self._states[0] + + class ShouldRouteToTdlibTests(unittest.TestCase): def test_routes_when_all_conditions_met(self): self.assertTrue( @@ -74,6 +104,91 @@ def test_rejects_unknown_size(self): ) +class StartClientReadyTests(unittest.IsolatedAsyncioTestCase): + async def test_returns_when_ready(self): + client = _FakeClient(["authorizationStateReady"]) + await start_client_ready(client, timeout_seconds=1.0, poll_interval_seconds=0.05) + self.assertEqual(client.start_calls, [False]) # start(wait_login=False) + + async def test_returns_once_state_becomes_ready(self): + client = _FakeClient( + [ + "authorizationStateWaitTdlibParameters", + "authorizationStateWaitTdlibParameters", + "authorizationStateReady", + ] + ) + await start_client_ready(client, timeout_seconds=1.0, poll_interval_seconds=0.01) + + async def test_raises_instead_of_hanging_when_unauthorized(self): + client = _FakeClient(["authorizationStateWaitPhoneNumber"]) + with self.assertRaisesRegex(TdlibDownloadError, "not ready"): + await start_client_ready( + client, timeout_seconds=0.2, poll_interval_seconds=0.05 + ) + + +class AssertExpectedSizeTests(unittest.TestCase): + def test_no_expected_size_skips_check(self): + assert_expected_size(downloaded_bytes=123, expected_bytes=None) # no raise + + def test_matching_size_passes(self): + assert_expected_size(downloaded_bytes=52_428_800, expected_bytes=52_428_800) + + def test_mismatch_raises_so_caller_falls_back(self): + with self.assertRaisesRegex(TdlibDownloadError, "different message"): + assert_expected_size(downloaded_bytes=10, expected_bytes=52_428_800) + + +class _FakeDownloadClient: + """Async fake exercising the progress-poll loop without pytdbot/network. + + getFile reports ``downloaded_size`` growing across polls until it reaches + ``total`` and flips ``is_downloading_completed``. Returns plain dicts, which + ``raise_if_error`` passes through unchanged (they are not pytdbot Errors). + """ + + def __init__(self, total, steps): + self.total = total + self._steps = list(steps) + self.download_calls = [] + + async def downloadFile(self, **kwargs): + self.download_calls.append(kwargs) + return {"local": {"downloaded_size": 0, "is_downloading_completed": False}} + + async def getFile(self, *, file_id): + downloaded = self._steps.pop(0) if len(self._steps) > 1 else self._steps[0] + return { + "size": self.total, + "expected_size": self.total, + "local": { + "downloaded_size": downloaded, + "is_downloading_completed": downloaded >= self.total, + }, + } + + +@unittest.skipUnless(PYTDBOT_AVAILABLE, "pytdbot not installed (optional [tdlib] extra)") +class DownloadWithProgressTests(unittest.IsolatedAsyncioTestCase): + async def test_reports_progress_and_returns_completed_file(self): + from telegram_mcp.tdlib_download import _download_with_progress + + client = _FakeDownloadClient(total=100, steps=[40, 100]) + seen = [] + result = await _download_with_progress( + client, + file_id=7, + progress_callback=lambda done, total: seen.append((done, total)), + total_bytes=100, + poll_interval_seconds=0, + ) + self.assertEqual(client.download_calls[0]["synchronous"], False) + self.assertIn((40, 100), seen) + self.assertIn((100, 100), seen) + self.assertTrue(result["local"]["is_downloading_completed"]) + + class AssertIsolatedFromTelethonTests(unittest.TestCase): def test_rejects_telethon_session_tree(self): with self.assertRaises(ValueError): @@ -111,6 +226,20 @@ def test_extract_file_id_from_message_video(self): ) self.assertEqual(extract_file_id_from_message(message), 555) + def test_extract_file_id_from_message_animation(self): + import pytdbot + + from telegram_mcp.tdlib_download import extract_file_id_from_message + + message = pytdbot.types.Message( + content=pytdbot.types.MessageAnimation( + animation=pytdbot.types.Animation( + animation=pytdbot.types.File(id=777, size=52_428_800) + ) + ) + ) + self.assertEqual(extract_file_id_from_message(message), 777) + def test_extract_file_id_from_message_unsupported_type(self): import pytdbot