feat(api,mcp): explain account_id mixed up with a telegram user id
This commit is contained in:
1 parent
676680319b
commit
ab59ef4525
2 files changed
+56
-1
No files matched your search
@@ -4,13 +4,17 @@ 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
|
||||
from pydantic import BaseModel
|
||||
|
||||
from api.mcp.format import build_transcript, load_notes, load_photos, resolve_names
|
||||
from dependencies.container import container
|
||||
from utils.jobs import enqueue
|
||||
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.search.models import SearchFilters
|
||||
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:
|
||||
- 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,
|
||||
@@ -63,6 +70,25 @@ 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]
|
||||
|
||||
|
||||
Reference in new issue
Block a user