feat(api,mcp): explain account_id mixed up with a telegram user id
This commit is contained in:
@@ -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]
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user