From 283eab1eb0c57b21d1a820f6570934633e24e284 Mon Sep 17 00:00:00 2001 From: h Date: Thu, 8 Oct 2026 21:22:58 +0200 Subject: [PATCH] feat(mcp): download any message file via share link --- backend/src/api/mcp/download.py | 113 ++++++++++++++++++++++++++++++++ backend/src/api/mcp/server.py | 68 +++++++++++++------ 2 files changed, 162 insertions(+), 19 deletions(-) create mode 100644 backend/src/api/mcp/download.py diff --git a/backend/src/api/mcp/download.py b/backend/src/api/mcp/download.py new file mode 100644 index 0000000..d7d4a26 --- /dev/null +++ b/backend/src/api/mcp/download.py @@ -0,0 +1,113 @@ +import asyncio +from pathlib import Path +from typing import Any + +import asyncpg +from fastmcp.exceptions import ToolError +from fastmcp.server.dependencies import get_http_request +from fastmcp.utilities.types import Image +from mcp.types import TextContent + +from utils.files import expiry_from_seconds +from utils.jobs import enqueue +from utils.read import chats, media, shares +from utils.read.models import FileShareView, MediaView +from utils.storage import ContentAddressedStorage + +LINK_TTL_SECONDS = 24 * 3600 +MAX_WAIT_SECONDS = 120 +_POLL_SECONDS = 1.0 +_INLINE_IMAGE_BYTES = 5 * 1024 * 1024 +_INLINE_TEXT_BYTES = 256 * 1024 +_INLINE_IMAGE_MIMES = {"image/jpeg", "image/png", "image/gif", "image/webp"} +_TEXT_MIMES = {"application/json", "application/xml", "application/x-yaml"} +_FAILED_JOB = ("failed", "canceled") + + +async def _job_error(pool: asyncpg.Pool, job_id: int) -> str | None: + row = await pool.fetchrow("SELECT status, error FROM jobs WHERE id = $1", job_id) + if row is None or row["status"] in _FAILED_JOB: + return (row and row["error"]) or "download job is gone" + return None + + +async def fetch_media( + pool: asyncpg.Pool, + account_id: int, + chat_id: int, + message_id: int, + wait_seconds: float, +) -> MediaView | None: + item = await media.get_message_media(pool, account_id, chat_id, message_id) + if item is None: + msg = f"message {chat_id}/{message_id} has no media" + raise ToolError(msg) + if item.downloaded and item.storage_key: + return item + job_id = await enqueue( + pool, account_id, "fetch_media", {"chat_id": chat_id, "message_id": message_id} + ) + loop = asyncio.get_running_loop() + deadline = loop.time() + wait_seconds + while loop.time() < deadline: + await asyncio.sleep(_POLL_SECONDS) + item = await media.get_message_media(pool, account_id, chat_id, message_id) + if item is not None and item.downloaded and item.storage_key: + return item + error = await _job_error(pool, job_id) + if error is not None: + msg = f"download failed: {error}" + raise ToolError(msg) + return None + + +def _origin() -> str: + request = get_http_request() + proto = request.headers.get("x-forwarded-proto", request.url.scheme) + return f"{proto}://{request.url.netloc}" + + +async def share_media(pool: asyncpg.Pool, item: MediaView, mime: str) -> FileShareView: + storage_key = item.storage_key or "" + share = await shares.find_active_share(pool, item.account_id, storage_key) + if share is None: + chat = await chats.get_chat(pool, item.account_id, item.chat_id) + share = await shares.create_share( + pool, + item.account_id, + "media", + storage_key, + item.file_name or f"media_{item.id}", + mime=mime, + file_size=item.file_size, + title=chat.title if chat else None, + chat_id=item.chat_id, + message_id=item.message_id, + expires_at=expiry_from_seconds(LINK_TTL_SECONDS), + ) + return share + + +def _is_text(mime: str) -> bool: + return mime.startswith("text/") or mime in _TEXT_MIMES + + +def inline_content(storage: ContentAddressedStorage, key: str, mime: str) -> list[Any]: + size = Path(storage.url(key)).stat().st_size + if mime in _INLINE_IMAGE_MIMES and size <= _INLINE_IMAGE_BYTES: + return [Image(data=storage.get(key), format=mime.rsplit("/", maxsplit=1)[-1])] + if _is_text(mime) and size <= _INLINE_TEXT_BYTES: + text = storage.get(key).decode("utf-8", errors="replace") + return [TextContent(type="text", text=text)] + return [] + + +def describe(item: MediaView, mime: str, share: FileShareView) -> str: + size = f"{item.file_size} bytes" if item.file_size else "size unknown" + expires = share.expires_at.isoformat() if share.expires_at else "never" + return ( + f"{item.file_name} ({item.kind}, {mime}, {size})\n" + f"Download: {_origin()}/f/{share.token}/{share.url_name}\n" + f"Public link, expires {expires}; fetch it with curl/wget or any HTTP " + "client, no auth needed." + ) diff --git a/backend/src/api/mcp/server.py b/backend/src/api/mcp/server.py index f810057..2ad9f7c 100644 --- a/backend/src/api/mcp/server.py +++ b/backend/src/api/mcp/server.py @@ -7,12 +7,19 @@ from fastmcp import FastMCP from fastmcp.exceptions import ToolError from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext from fastmcp.tools.base import ToolResult -from mcp.types import CallToolRequestParams +from mcp.types import CallToolRequestParams, TextContent from pydantic import BaseModel +from api.mcp.download import ( + MAX_WAIT_SECONDS, + describe, + fetch_media, + inline_content, + share_media, +) from api.mcp.format import build_transcript, load_notes, load_photos, resolve_names from dependencies.container import container -from utils.jobs import enqueue +from utils.files import resolve_mime from utils.read import annotations, chats, media, peers, presence, social, watches from utils.read.accounts import self_user_id, unknown_account_hint from utils.read.models import DEFAULT_LIMIT, Page @@ -44,7 +51,7 @@ Typical flows: `before_id` (first id you saw), to newer with `after_id` (last id you saw); begin a full forward walk with `after_id=0`. Each transcript line starts with its `#message_id`, and photos show as `[photo #message_id]` — pass that - id to `get_media` (with `fetch=True` to download) or to annotation tools. + id to `download_media` or to annotation tools. Lines marked "📝 [your private note, NOT in Telegram]" are the user's own annotations attached locally via the web UI; they never existed in Telegram. Treat them @@ -53,8 +60,11 @@ Typical flows: transcripts; supports chat/sender/date filters and regex). - Forensics: `get_deleted_messages` and `get_message_versions` recover content removed or edited in Telegram but kept in the archive. -- Media: `get_media` returns metadata; pass `fetch=True` to enqueue a lazy - download if the file isn't stored yet. +- Files: `download_media` gets any message attachment (photo, video, voice, + document, ...). It pulls the file from Telegram if the archive lacks it and + returns a public link valid 24h that needs no auth (`curl -O `), plus + the content itself for small images and text files. `get_media` returns + metadata only. - Monitoring: `set_watch` creates a local rule, `list_watches` lists rules, and `list_alerts` reads what those rules fired. @@ -63,7 +73,8 @@ Watched bots: own `account_id` and its chats, messages and media read through the same tools as a user account. Nothing bot-specific to call. -Writes: everything is read-only except `set_watch`, the only allowed write. +Writes: everything is read-only except `set_watch` and the share links +`download_media` creates (visible and revocable in the web UI). Paging: list/search tools accept `limit` and `offset`. """ @@ -239,21 +250,40 @@ async def get_message_versions( @mcp.tool async def get_media( - account_id: int, chat_id: int, message_id: int, fetch: bool = False + account_id: int, chat_id: int, message_id: int ) -> dict[str, Any] | None: - """Get media metadata for a message; set fetch=True to enqueue lazy download.""" + """Get media metadata for a message; use download_media to get the file.""" + item = await media.get_message_media(await _pool(), account_id, chat_id, message_id) + return item.model_dump(mode="json") if item else None + + +@mcp.tool(output_schema=None) +async def download_media( + account_id: int, chat_id: int, message_id: int, wait_seconds: int = 45 +) -> list[Any]: + """Download the file attached to a message: any photo, video, voice, document. + + Pulls it from Telegram first if the archive doesn't have it yet, waiting up + to `wait_seconds`. Returns a public download link (valid 24h, no auth) and + inlines small images and text files. + """ pool = await _pool() - item = await media.get_message_media(pool, account_id, chat_id, message_id) - if item is None: - return None - if fetch and not item.downloaded: - await enqueue( - pool, - account_id, - "fetch_media", - {"chat_id": chat_id, "message_id": message_id}, - ) - return item.model_dump(mode="json") + wait = min(max(wait_seconds, 0), MAX_WAIT_SECONDS) + item = await fetch_media(pool, account_id, chat_id, message_id, wait) + if item is None or item.storage_key is None: + return [ + TextContent( + type="text", + text="Still downloading from Telegram; call download_media again.", + ) + ] + mime = resolve_mime(item.kind, item.mime, item.file_name) + share = await share_media(pool, item, mime) + storage = await container.get(ContentAddressedStorage) + return [ + TextContent(type="text", text=describe(item, mime, share)), + *inline_content(storage, item.storage_key, mime), + ] @mcp.tool