feat(api,mcp): explain account_id mixed up with a telegram user id

This commit is contained in:
hh
2026-09-06 21:30:51 +02:00
parent 676680319b
commit ab59ef4525
2 changed files with 56 additions and 1 deletions
+27 -1
View File
@@ -4,13 +4,17 @@ from typing import Any
import asyncpg import asyncpg
from fastmcp import FastMCP 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 pydantic import BaseModel from pydantic import BaseModel
from api.mcp.format import build_transcript, load_notes, load_photos, resolve_names from api.mcp.format import build_transcript, load_notes, load_photos, resolve_names
from dependencies.container import container from dependencies.container import container
from utils.jobs import enqueue from utils.jobs import enqueue
from utils.read import annotations, chats, media, peers, presence, social, watches from utils.read import annotations, chats, media, peers, presence, social, watches
from utils.read.accounts import self_user_id from utils.read.accounts import self_user_id, unknown_account_hint
from utils.read.models import DEFAULT_LIMIT, Page from utils.read.models import DEFAULT_LIMIT, Page
from utils.search.models import SearchFilters from utils.search.models import SearchFilters
from utils.search.repository import search_messages from utils.search.repository import search_messages
@@ -23,6 +27,9 @@ peer history) into Postgres and exposes it read-only over these tools.
Account scoping: Account scoping:
- Every tool takes `account_id`. It selects which archived Telegram account to - 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`. 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: Identifiers:
- `chat_id` and `peer_id` are Telegram IDs (negative for groups/channels, - `chat_id` and `peer_id` are Telegram IDs (negative for groups/channels,
@@ -63,6 +70,25 @@ async def _pool() -> asyncpg.Pool:
return await container.get(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]]: def _dump(items: Sequence[BaseModel]) -> list[dict[str, Any]]:
return [item.model_dump(mode="json") for item in items] return [item.model_dump(mode="json") for item in items]
+29
View File
@@ -45,6 +45,35 @@ async def self_user_id(pool: asyncpg.Pool, account_id: int) -> int | None:
return tg_user_id return tg_user_id
_known_ids: set[int] = set()
def _account_label(row: asyncpg.Record) -> str:
named = f", {row['label']}" if row["label"] else ""
return f"account_id={row['account_id']} (Telegram id {row['tg_user_id']}{named})"
async def unknown_account_hint(pool: asyncpg.Pool, account_id: int) -> str | None:
"""Explain an account_id that matches no account, or None if it is valid."""
if account_id in _known_ids:
return None
rows = await pool.fetch(
"SELECT account_id, tg_user_id, label FROM accounts ORDER BY account_id"
)
_known_ids.update(row["account_id"] for row in rows)
if account_id in _known_ids:
return None
known = "; ".join(_account_label(row) for row in rows) or "none"
mistaken = next((row for row in rows if row["tg_user_id"] == account_id), None)
if mistaken is not None:
return (
f"{account_id} is a Telegram user id, not an account_id. In beavergram "
f"that account is account_id={mistaken['account_id']}. "
f"Retry with account_id={mistaken['account_id']}."
)
return f"No account with account_id={account_id}. Archived accounts: {known}."
async def list_accounts(pool: asyncpg.Pool) -> list[AccountView]: async def list_accounts(pool: asyncpg.Pool) -> list[AccountView]:
rows = await pool.fetch( rows = await pool.fetch(
f"SELECT {_ACCOUNT_COLS} FROM accounts ORDER BY account_id" # noqa: S608 f"SELECT {_ACCOUNT_COLS} FROM accounts ORDER BY account_id" # noqa: S608