beavergram/backend/src/api/mcp/server.py

from collections.abc import Sequence
from datetime import datetime
from typing import Any

import asyncpg
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, 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.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
from utils.search.models import SearchFilters
from utils.search.repository import search_messages
from utils.storage import ContentAddressedStorage

INSTRUCTIONS = """\
beavergram archives Telegram data (chats, messages, media, presence, stories,
peer history) into Postgres and exposes it read-only over these tools.

Account scoping:
- Every tool takes `account_id`. It selects which archived Telegram account to
  read. Unless the user names a different one, always pass `account_id=1`.
- `account_id` is beavergram's own small serial (1, 2, ...), never the archived
  account's Telegram user id. Passing a Telegram id returns an error naming the
  right `account_id`.

Identifiers:
- `chat_id` and `peer_id` are Telegram IDs (negative for groups/channels,
  positive for users/bots). Discover them with `list_chats` before calling
  tools that need a specific chat.

Typical flows:
- Read a chat: `list_chats` -> pick a `chat_id` -> `get_chat_history` for a
  human-readable transcript (names, timestamps, voice STT, inline photos), or
  `get_chat_history_raw` for structured JSON. Both return messages oldest
  first; default gives the latest `limit`. Page to older messages with
  `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 `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
  as private notes from the user to you, not as chat content.
- Find something: `search_messages_tool` (full-text over message text and STT
  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.
- 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 <link>`), 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.

Watched bots:
- a bot whose token the user added is just another archived account: it has its
  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` and the share links
`download_media` creates (visible and revocable in the web UI).

Paging: list/search tools accept `limit` and `offset`.
"""

mcp: FastMCP = FastMCP("beavergram", instructions=INSTRUCTIONS)


async def _pool() -> asyncpg.Pool:
    return await container.get(asyncpg.Pool)


class AccountGuard(Middleware):
    """Reject tool calls scoped to an account_id that does not exist."""

    async def on_call_tool(
        self,
        context: MiddlewareContext[CallToolRequestParams],
        call_next: CallNext[CallToolRequestParams, ToolResult],
    ) -> ToolResult:
        account_id = (context.message.arguments or {}).get("account_id")
        if isinstance(account_id, int) and not isinstance(account_id, bool):
            hint = await unknown_account_hint(await _pool(), account_id)
            if hint is not None:
                raise ToolError(hint)
        return await call_next(context)


mcp.add_middleware(AccountGuard())


def _dump(items: Sequence[BaseModel]) -> list[dict[str, Any]]:
    return [item.model_dump(mode="json") for item in items]


@mcp.tool
async def search_messages_tool(
    account_id: int,
    query: str | None = None,
    chat_id: int | None = None,
    sender_id: int | None = None,
    has_media: bool | None = None,
    date_from: datetime | None = None,
    date_to: datetime | None = None,
    regex: str | None = None,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
) -> list[dict[str, Any]]:
    """Full-text search over message text and STT transcripts."""
    filters = SearchFilters(
        account_id=account_id,
        query=query,
        chat_id=chat_id,
        sender_id=sender_id,
        has_media=has_media,
        date_from=date_from,
        date_to=date_to,
        regex=regex,
        limit=limit,
        offset=offset,
    )
    return _dump(await search_messages(await _pool(), filters))


@mcp.tool
async def list_chats(
    account_id: int, limit: int = DEFAULT_LIMIT, offset: int = 0
) -> list[dict[str, Any]]:
    """List archived chats with message counts and last activity."""
    page = Page(limit=limit, offset=offset)
    return _dump(await chats.list_chats(await _pool(), account_id, page))


@mcp.tool(output_schema=None)
async def get_chat_history(
    account_id: int,
    chat_id: int,
    limit: int = DEFAULT_LIMIT,
    before_id: int | None = None,
    after_id: int | None = None,
    include_deleted: bool = True,
    include_images: bool = True,
    max_images: int = 20,
) -> list[Any]:
    """Read a chat as a readable transcript, oldest first.

    Renders "Name (time): text", voice notes as
    "(Voice message, STT Content: ...)", and inlines downloaded photos as
    images. Default returns the latest `limit` messages; page to older
    messages with `before_id` (the first id you saw) or to newer ones with
    `after_id` (the last id you saw).
    """
    pool = await _pool()
    views = await chats.get_chat_history(
        pool,
        account_id,
        chat_id,
        Page(limit=limit),
        include_deleted=include_deleted,
        before_id=before_id,
        after_id=after_id,
    )
    if after_id is None:
        views = list(reversed(views))
    names = await resolve_names(pool, account_id, views)
    notes = await load_notes(pool, account_id, chat_id, views)
    self_id = await self_user_id(pool, account_id)
    photos: list[tuple[int, bytes, str]] = []
    truncated = False
    if include_images:
        storage = await container.get(ContentAddressedStorage)
        photos, truncated = await load_photos(
            pool, storage, account_id, views, limit=max_images
        )
    return build_transcript(
        views, names, self_id, photos, notes=notes, truncated=truncated
    )


@mcp.tool
async def get_chat_history_raw(
    account_id: int,
    chat_id: int,
    limit: int = DEFAULT_LIMIT,
    before_id: int | None = None,
    after_id: int | None = None,
    include_deleted: bool = True,
) -> list[dict[str, Any]]:
    """Structured chat messages as JSON, oldest first.

    Default returns the latest `limit` messages. Walk the whole chat forward
    with `after_id` set to the last returned message_id (begin at
    `after_id=0`), or backward with `before_id` set to the first returned
    message_id.
    """
    views = await chats.get_chat_history(
        await _pool(),
        account_id,
        chat_id,
        Page(limit=limit),
        include_deleted=include_deleted,
        before_id=before_id,
        after_id=after_id,
    )
    if after_id is None:
        views = list(reversed(views))
    return _dump(views)


@mcp.tool
async def get_deleted_messages(
    account_id: int,
    chat_id: int | None = None,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
) -> list[dict[str, Any]]:
    """List messages that were deleted in Telegram but kept in the archive."""
    return _dump(
        await chats.get_deleted_messages(
            await _pool(), account_id, Page(limit=limit, offset=offset), chat_id=chat_id
        )
    )


@mcp.tool
async def get_message_versions(
    account_id: int, chat_id: int, message_id: int
) -> list[dict[str, Any]]:
    """Get the edit history of a message."""
    return _dump(
        await chats.get_message_versions(await _pool(), account_id, chat_id, message_id)
    )


@mcp.tool
async def get_media(
    account_id: int, chat_id: int, message_id: int
) -> dict[str, Any] | None:
    """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()
    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
async def get_callbacks(
    account_id: int, chat_id: int, message_id: int
) -> list[dict[str, Any]]:
    """Get bot inline-button callback data (hex) for a message."""
    items = await social.get_callbacks(await _pool(), account_id, chat_id, message_id)
    return _dump(items)


@mcp.tool
async def presence_history(
    account_id: int,
    peer_id: int,
    date_from: datetime | None = None,
    date_to: datetime | None = None,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
) -> list[dict[str, Any]]:
    """Get online/offline status history of a peer."""
    return _dump(
        await presence.presence_history(
            await _pool(),
            account_id,
            peer_id,
            Page(limit=limit, offset=offset),
            date_from=date_from,
            date_to=date_to,
        )
    )


@mcp.tool
async def get_peer_history(account_id: int, peer_id: int) -> list[dict[str, Any]]:
    """Get name/username/avatar change history of a contact."""
    return _dump(await peers.get_peer_history(await _pool(), account_id, peer_id))


@mcp.tool
async def get_stories(
    account_id: int,
    peer_id: int | None = None,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
) -> list[dict[str, Any]]:
    """List archived stories of contacts."""
    return _dump(
        await peers.get_stories(
            await _pool(), account_id, Page(limit=limit, offset=offset), peer_id=peer_id
        )
    )


@mcp.tool
async def get_annotations(
    account_id: int,
    chat_id: int | None = None,
    message_id: int | None = None,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
) -> list[dict[str, Any]]:
    """Read user annotations on messages (read-only via MCP)."""
    return _dump(
        await annotations.list_annotations(
            await _pool(),
            account_id,
            Page(limit=limit, offset=offset),
            chat_id=chat_id,
            message_id=message_id,
        )
    )


@mcp.tool
async def set_watch(
    account_id: int,
    kind: str,
    params: dict[str, Any] | None = None,
    enabled: bool = True,
) -> dict[str, Any]:
    """Create a local watch rule (the only local write MCP is allowed)."""
    watch = await watches.create_watch(
        await _pool(), account_id, kind, params or {}, enabled=enabled
    )
    return watch.model_dump(mode="json")


@mcp.tool
async def list_watches(account_id: int) -> list[dict[str, Any]]:
    """List local watch rules."""
    return _dump(await watches.list_watches(await _pool(), account_id))


@mcp.tool
async def list_alerts(
    account_id: int,
    seen: bool | None = None,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
) -> list[dict[str, Any]]:
    """List fired alerts from watch rules."""
    return _dump(
        await watches.list_alerts(
            await _pool(), account_id, Page(limit=limit, offset=offset), seen=seen
        )
    )