feat(mcp): download any message file via share link

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."
+    )