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,
    people,
    presence,
    presence_stats,
    social,
    watches,
)
from utils.read.accounts import self_user_id, unknown_account_hint
from utils.read.models import DEFAULT_LIMIT, Page
from utils.read.period import DEFAULT_TZ, MAX_DAYS, known_timezone
from utils.read.read_receipts import read_info
from utils.search.models import SearchFilters
from utils.search.repository import search_messages
from utils.storage import ContentAddressedStorage
from utils.watch_rules import WatchKind, WatchParamsError, clean_name, normalize_params

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).
- People: `get_person` gives one person in full: profile changes, shared groups,
  joins and leaves, deleted message count, the private chat.
- 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 (kinds `keyword`,
  `new_message`, `message_deleted`, `message_edited`, `peer_online`; params in
  its description), `list_watches` lists rules, and `list_alerts` reads what
  those rules fired, with chat, message preview and a web client `path`.

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 get_read_info(
    account_id: int, chat_id: int, message_id: int
) -> dict[str, Any] | None:
    """When the user's own message was read, from the archive only.

    `observed_at` is when the archive saw the chat read up to this message,
    `exact_at` the read date Telegram reported for a private chat, `readers`
    who read it in a group and when, `listened_at` when a voice or round
    video was played. `outcome` is the last Telegram answer: ok, not_read_yet,
    user_privacy (they hide read time), your_privacy (the user hides theirs),
    chat_too_big, too_old, unavailable, flood or error.
    """
    info = await read_info(await _pool(), account_id, chat_id, message_id)
    return info.model_dump(mode="json") if info else None


@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(name="presence_stats")
async def presence_stats_tool(
    account_id: int,
    peer_id: int,
    days: int = 30,
    tz: str = DEFAULT_TZ,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
) -> dict[str, Any]:
    """Time a peer spent online, from status pushes turned into online sessions.

    `tz` is an IANA zone name; days and hours are counted in it. `week` holds
    the average online minutes for each weekday (0 is Monday) and hour, null
    where the archive has no data for that hour. `days` gives online seconds
    per day, `gaps` the stretches when capture was down, `sessions` the latest
    sessions (online stretches with breaks under 5 minutes), paged by
    `limit`/`offset`. `hidden` means the peer hides their last seen time.
    """
    pool = await _pool()
    if not 1 <= days <= MAX_DAYS:
        msg = f"days must be between 1 and {MAX_DAYS}"
        raise ToolError(msg)
    if not await known_timezone(pool, tz):
        msg = f"Unknown time zone: {tz}"
        raise ToolError(msg)
    stats = await presence_stats.presence_stats(
        pool, account_id, peer_id, Page(limit=limit, offset=offset), days=days, tz=tz
    )
    return stats.model_dump(mode="json")


@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_person(account_id: int, peer_id: int) -> dict[str, Any]:
    """Everything the archive knows about one person.

    Current profile, name/username/phone changes with times (newest first),
    avatar history, groups shared in the archive with their message counts and
    last message, joins and leaves with who added or removed them, the number
    of their deleted messages, and the private chat if there is one.
    """
    pool = await _pool()
    overview = await people.get_overview(pool, account_id, peer_id)
    groups = await people.get_groups(pool, account_id, peer_id)
    return {**overview.model_dump(mode="json"), **groups.model_dump(mode="json")}


@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: WatchKind,
    params: dict[str, Any] | None = None,
    name: str | None = None,
    enabled: bool = True,
) -> dict[str, Any]:
    """Create a local watch rule (the only local write MCP is allowed).

    A watch fires an alert into the web client (live counter and list) and to
    the user's webhooks. The user's own messages never fire. Kinds and params:

    - `keyword`: `{"pattern": str, "regex": bool = false, "chat_id"?: int,
      "sender_id"?: int}`. Case-insensitive substring, or a Python regex when
      `regex` is true (at most 200 characters, no nested repeats like
      `(a+)+`). Checked on new messages, on edits (the new text) and on voice
      transcripts. Without `chat_id` it watches every chat.
    - `new_message`: `{"chat_id"?: int, "sender_id"?: int}`, at least one: any
      new message in that chat, or from that person anywhere, or both.
    - `message_deleted`: `{"chat_id"?: int}`. A stored message deleted in that
      chat; without `chat_id`, in any private chat.
    - `message_edited`: `{"chat_id"?: int}`. Same scopes as `message_deleted`.
    - `peer_online`: `{"peer_id": int}`. The person comes online after at
      least 15 minutes away.

    `name` is an optional human label shown in the list and in alerts.
    """
    try:
        clean = normalize_params(kind, params or {})
    except WatchParamsError as exc:
        raise ToolError(str(exc)) from exc
    watch = await watches.create_watch(
        await _pool(), account_id, kind, clean, name=clean_name(name), 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,
    before: int | None = None,
    limit: int = DEFAULT_LIMIT,
) -> dict[str, Any]:
    """List fired alerts from watch rules, newest first.

    Each alert names its watch (id, kind, name), the event (`message`, `edit`,
    `transcript`, `delete`, `online`), the chat, the message preview with its
    sender and deleted/edited flags, the peer for `peer_online`, and `path`,
    the web client link. `seen=false` keeps unseen ones. Page older with
    `before` set to `next_before` from the previous call; `unseen` is the
    account's unseen count.
    """
    page = await watches.list_alerts(
        await _pool(), account_id, seen=seen, before=before, limit=limit
    )
    return page.model_dump(mode="json")