From b9cf9452f4fae5d3fbd77c20ffd87e84a120c6d4 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Fri, 30 Jan 2026 08:39:27 +0530 Subject: [PATCH] Enhance MCP tool classification and API fallback mechanisms - Introduced a three-tier tool classification system for rate limiting: state_only, api_read, and api_write, improving the management of tool execution based on their API call characteristics. - Updated rate limiting configuration to reflect new classifications, ensuring appropriate delays for different tool types. - Implemented API fallback mechanisms in various tools to first attempt data retrieval from cached Redux state before falling back to Telegram API calls, enhancing performance and reliability. - Refactored multiple Telegram tools to utilize the new API fallback approach, improving user experience during data retrieval. This update significantly enhances the efficiency and responsiveness of the MCP by optimizing tool execution and data retrieval processes. --- src/lib/mcp/index.ts | 3 + src/lib/mcp/rateLimiter.ts | 227 +++++++++----- src/lib/mcp/telegram/telegramApi.ts | 284 ++++++++++++++++++ src/lib/mcp/telegram/tools/getChat.ts | 113 +++++-- src/lib/mcp/telegram/tools/getChats.ts | 4 +- src/lib/mcp/telegram/tools/getHistory.ts | 4 +- src/lib/mcp/telegram/tools/getMe.ts | 4 +- .../mcp/telegram/tools/getMessageContext.ts | 4 +- src/lib/mcp/telegram/tools/getMessages.ts | 4 +- src/lib/mcp/telegram/tools/listChats.ts | 4 +- src/lib/mcp/telegram/tools/listMessages.ts | 4 +- src/lib/mcp/telegram/tools/resolveUsername.ts | 84 +++++- src/lib/mcp/telegram/tools/searchMessages.ts | 73 ++++- .../mcp/telegram/tools/searchPublicChats.ts | 99 +++++- 14 files changed, 797 insertions(+), 114 deletions(-) diff --git a/src/lib/mcp/index.ts b/src/lib/mcp/index.ts index 8a330345e..44578fcd3 100644 --- a/src/lib/mcp/index.ts +++ b/src/lib/mcp/index.ts @@ -11,8 +11,11 @@ export { SocketIOMCPTransportImpl } from "./transport"; export { enforceRateLimit, resetRequestCallCount, + classifyTool, + isStateOnlyTool, isReadOnlyTool, isHeavyTool, getRateLimitStatus, RATE_LIMIT_CONFIG, } from "./rateLimiter"; +export type { ToolTier } from "./rateLimiter"; diff --git a/src/lib/mcp/rateLimiter.ts b/src/lib/mcp/rateLimiter.ts index d4ca3f608..958d491f1 100644 --- a/src/lib/mcp/rateLimiter.ts +++ b/src/lib/mcp/rateLimiter.ts @@ -1,13 +1,14 @@ /** * MCP Rate Limiter * - * Three-tier rate limiting system for MCP tool execution: - * 1. Per-request counter — caps tool calls within a single agent request - * 2. Per-minute sliding window — prevents sustained high-frequency usage - * 3. Inter-call delay — enforces minimum gap between API calls + * Three-tier tool classification: + * 1. STATE_ONLY — reads cached Redux state, zero API calls → no rate limit + * 2. API_READ — reads from Telegram API → standard inter-call delay + * 3. API_WRITE — mutates state on Telegram servers → heavy inter-call delay * - * Read-only tools that only access cached Redux state bypass rate limits. - * Mutation tools (send, delete, create, etc.) incur heavier delays. + * On top of the per-call delay, two budget caps apply to all API-bound tools: + * - Per-request counter (caps tool calls within a single agent request) + * - Per-minute sliding window (prevents sustained high-frequency usage) */ import { mcpLog, mcpWarn } from "./logger"; @@ -17,13 +18,13 @@ import { mcpLog, mcpWarn } from "./logger"; // --------------------------------------------------------------------------- export const RATE_LIMIT_CONFIG = { - /** Minimum delay (ms) between ANY tool calls that hit the Telegram API */ - MIN_CALL_DELAY_MS: 500, - /** Extra delay (ms) for mutation/write operations */ - HEAVY_CALL_DELAY_MS: 1000, - /** Maximum tool calls allowed within a 60-second sliding window */ + /** Minimum delay (ms) between API read calls */ + API_READ_DELAY_MS: 500, + /** Delay (ms) between API write / mutation calls */ + API_WRITE_DELAY_MS: 1000, + /** Maximum API-bound tool calls within a 60-second sliding window */ MAX_CALLS_PER_MINUTE: 30, - /** Maximum tool calls allowed within a single MCP request */ + /** Maximum API-bound tool calls within a single MCP request */ MAX_CALLS_PER_REQUEST: 20, } as const; @@ -31,63 +32,137 @@ export const RATE_LIMIT_CONFIG = { // Tool classification // --------------------------------------------------------------------------- +export type ToolTier = "state_only" | "api_read" | "api_write"; + /** - * Read-only tools that only query cached Redux state — zero Telegram API calls. - * These are exempt from rate limiting entirely. + * Tools that ONLY read from cached Redux state — zero Telegram API calls. + * Bypass all rate limiting; execute instantly. */ -const READ_ONLY_TOOLS = new Set([ +const STATE_ONLY_TOOLS = new Set([ + // Chat state (selectOrderedChats / state.chats) "get_chats", "list_chats", "get_chat", + + // Message state (state.messages / state.messagesOrder) "get_messages", "list_messages", - "get_me", - "get_drafts", "get_message_context", - "get_pinned_messages", "get_history", - "get_user_status", - "get_participants", - "get_admins", - "get_blocked_users", - "list_contacts", - "search_contacts", - "get_contact_ids", + + // Current user (state.currentUser) + "get_me", + + // Derived from cached chat/message data + "list_inline_buttons", "get_contact_chats", "get_direct_chat_by_contact", - "list_inline_buttons", - "list_topics", - "get_message_reactions", + + // These read from cached messages only (no API call) + "get_last_interaction", + "get_media_info", ]); /** - * Heavy (mutation) tools that modify state on Telegram servers. - * These incur HEAVY_CALL_DELAY_MS instead of MIN_CALL_DELAY_MS. + * Tools that call the Telegram API but only READ data (no mutations). + * Subject to standard inter-call delay + per-minute/per-request caps. */ -const HEAVY_TOOLS = new Set([ +const API_READ_TOOLS = new Set([ + // Contacts / users (contacts.GetContacts, contacts.Search, etc.) + "list_contacts", + "search_contacts", + "get_contact_ids", + "get_blocked_users", + "get_user_status", + "get_user_photos", + + // Chat metadata (channels.GetParticipants, messages.GetFullChat, etc.) + "get_participants", + "get_admins", + "get_banned_users", + "get_recent_actions", + "get_bot_info", + "get_privacy_settings", + + // Messages (messages.Search, messages.GetMessagesReactions, etc.) + "get_pinned_messages", + "get_message_reactions", + "search_messages", + + // Drafts / misc reads + "get_drafts", + "get_sticker_sets", + "get_gif_search", + + // Topics (channels.GetForumTopics) + "list_topics", + + // Discovery (these call the Telegram API for server-side search) + "search_public_chats", + "resolve_username", + "export_contacts", +]); + +/** + * Tools that MODIFY state on Telegram servers. + * Subject to heavy inter-call delay + per-minute/per-request caps. + */ +const API_WRITE_TOOLS = new Set([ + // Message mutations "send_message", "reply_to_message", - "forward_message", + "edit_message", "delete_message", + "forward_message", + "pin_message", + "unpin_message", + "mark_as_read", + "send_reaction", + "remove_reaction", + "save_draft", + "clear_draft", + "press_inline_button", + "create_poll", + + // Invite link (generates/exports a link — treated as write) + "get_invite_link", + + // Chat mutations "create_group", "create_channel", "invite_to_group", - "ban_user", - "unban_user", - "promote_admin", - "demote_admin", - "archive_chat", - "unarchive_chat", - "leave_chat", - "import_contacts", - "export_contacts", - "set_privacy_settings", - "set_bot_commands", - "set_profile_photo", - "delete_profile_photo", + "edit_chat_title", "edit_chat_photo", "delete_chat_photo", - "create_poll", + "leave_chat", + "archive_chat", + "unarchive_chat", + "mute_chat", + "unmute_chat", + "export_chat_invite", + "import_chat_invite", + "join_chat_by_link", + "subscribe_public_channel", + + // Admin / moderation + "promote_admin", + "demote_admin", + "ban_user", + "unban_user", + + // Contact mutations + "add_contact", + "delete_contact", + "block_user", + "unblock_user", + "import_contacts", + + // Profile mutations + "update_profile", + "set_profile_photo", + "delete_profile_photo", + "set_privacy_settings", + "set_bot_commands", ]); // --------------------------------------------------------------------------- @@ -108,23 +183,36 @@ const callHistory: number[] = []; // --------------------------------------------------------------------------- /** - * Returns true if the given tool name reads only from local cache and should - * bypass all rate limiting. + * Classify a tool into one of the three tiers. + * Unknown tools default to api_read (safe fallback — rate limited but not heavy). */ -export function isReadOnlyTool(toolName: string): boolean { - return READ_ONLY_TOOLS.has(toolName); +export function classifyTool(toolName: string): ToolTier { + if (STATE_ONLY_TOOLS.has(toolName)) return "state_only"; + if (API_WRITE_TOOLS.has(toolName)) return "api_write"; + if (API_READ_TOOLS.has(toolName)) return "api_read"; + // Unknown tools default to api_read so they're rate limited + return "api_read"; } /** - * Returns true if the given tool performs a mutation/write operation and should - * incur a heavier inter-call delay. + * Returns true if the tool only reads from local cache (no API call). + */ +export function isStateOnlyTool(toolName: string): boolean { + return STATE_ONLY_TOOLS.has(toolName); +} + +/** + * Returns true if the tool performs a mutation/write via the Telegram API. */ export function isHeavyTool(toolName: string): boolean { - return HEAVY_TOOLS.has(toolName); + return API_WRITE_TOOLS.has(toolName); } +/** @deprecated Use isStateOnlyTool instead */ +export const isReadOnlyTool = isStateOnlyTool; + /** - * Reset the per-request call counter. Call this at the start of each new + * Reset the per-request call counter. Call at the start of each new * MCP request (agent turn) to allow a fresh budget of tool calls. */ export function resetRequestCallCount(): void { @@ -134,18 +222,19 @@ export function resetRequestCallCount(): void { /** * Enforce rate limits before executing a tool. * - * - Read-only tools skip all limits. - * - For API-bound tools: + * - State-only tools skip all limits (instant). + * - API-bound tools (read or write): * 1. Check per-request budget → throw if exceeded * 2. Check per-minute sliding window → sleep until budget available - * 3. Enforce inter-call delay (heavier for mutation tools) + * 3. Enforce inter-call delay (500ms for reads, 1000ms for writes) * - * Call this BEFORE executing the tool handler. It may asynchronously wait - * if delays are needed, or throw if hard limits are exceeded. + * Call BEFORE executing the tool handler. May sleep or throw. */ export async function enforceRateLimit(toolName: string): Promise { - // Read-only tools are always allowed instantly - if (isReadOnlyTool(toolName)) { + const tier = classifyTool(toolName); + + // State-only tools are always allowed instantly + if (tier === "state_only") { return; } @@ -153,7 +242,7 @@ export async function enforceRateLimit(toolName: string): Promise { callsInCurrentRequest += 1; if (callsInCurrentRequest > RATE_LIMIT_CONFIG.MAX_CALLS_PER_REQUEST) { throw new Error( - `Rate limit: exceeded ${RATE_LIMIT_CONFIG.MAX_CALLS_PER_REQUEST} tool calls per request. ` + + `Rate limit: exceeded ${RATE_LIMIT_CONFIG.MAX_CALLS_PER_REQUEST} API tool calls per request. ` + `Try breaking your task into smaller steps.`, ); } @@ -163,7 +252,6 @@ export async function enforceRateLimit(toolName: string): Promise { purgeOldEntries(now); if (callHistory.length >= RATE_LIMIT_CONFIG.MAX_CALLS_PER_MINUTE) { - // Wait until the oldest entry expires from the window const oldestTimestamp = callHistory[0]; const waitMs = oldestTimestamp + 60_000 - now + 50; // +50ms buffer mcpWarn( @@ -174,15 +262,18 @@ export async function enforceRateLimit(toolName: string): Promise { purgeOldEntries(Date.now()); } - // --- Inter-call delay --- - const requiredDelay = isHeavyTool(toolName) - ? RATE_LIMIT_CONFIG.HEAVY_CALL_DELAY_MS - : RATE_LIMIT_CONFIG.MIN_CALL_DELAY_MS; + // --- Inter-call delay (tier-dependent) --- + const requiredDelay = + tier === "api_write" + ? RATE_LIMIT_CONFIG.API_WRITE_DELAY_MS + : RATE_LIMIT_CONFIG.API_READ_DELAY_MS; const elapsed = Date.now() - lastCallTime; if (elapsed < requiredDelay) { const waitMs = requiredDelay - elapsed; - mcpLog(`Rate limit: inter-call delay ${waitMs}ms for '${toolName}'`); + mcpLog( + `Rate limit: ${tier} delay ${waitMs}ms for '${toolName}'`, + ); await sleep(waitMs); } diff --git a/src/lib/mcp/telegram/telegramApi.ts b/src/lib/mcp/telegram/telegramApi.ts index 807457ed0..4da527918 100644 --- a/src/lib/mcp/telegram/telegramApi.ts +++ b/src/lib/mcp/telegram/telegramApi.ts @@ -1,6 +1,10 @@ /** * Telegram API helpers for MCP tools * Uses mtprotoService + Redux telegram state (alphahuman) + * + * "WithApiFallback" variants try cached Redux state first, then call the + * Telegram API when cache is empty. They call enforceRateLimit() internally + * before any API call so the fast cache path stays unthrottled. */ import { store } from "../../../store"; @@ -15,6 +19,9 @@ import type { TelegramUser, TelegramMessage, } from "../../../store/telegram/types"; +import { Api } from "telegram"; +import bigInt from "big-integer"; +import { enforceRateLimit } from "../rateLimiter"; function getTelegramState() { return selectTelegramUserState(store.getState()); @@ -202,3 +209,280 @@ export function formatMessage(message: TelegramMessage): FormattedMessage { } return result; } + +// --------------------------------------------------------------------------- +// API Fallback helpers +// --------------------------------------------------------------------------- + +/** Convert a raw GramJS message to our TelegramMessage format */ +function apiMessageToTelegramMessage( + msg: Record, + chatId: string, +): TelegramMessage { + const fromId = + msg.fromId && typeof msg.fromId === "object" && "userId" in (msg.fromId as object) + ? String((msg.fromId as { userId: unknown }).userId) + : undefined; + + const replyTo = msg.replyTo as { replyToMsgId?: number } | undefined; + + const media = msg.media as { className?: string } | undefined; + let mediaInfo: TelegramMessage["media"] | undefined; + if (media && media.className && media.className !== "MessageMediaEmpty") { + mediaInfo = { type: media.className }; + } + + return { + id: String(msg.id ?? ""), + chatId, + date: typeof msg.date === "number" ? msg.date : 0, + message: typeof msg.message === "string" ? msg.message : "", + fromId, + isOutgoing: Boolean(msg.out), + isEdited: msg.editDate != null, + isForwarded: msg.fwdFrom != null, + replyToMessageId: replyTo?.replyToMsgId != null + ? String(replyTo.replyToMsgId) + : undefined, + media: mediaInfo, + }; +} + +/** + * Get messages — tries cached Redux state first, falls back to Telegram API. + * Calls enforceRateLimit() internally before any API call. + */ +export async function getMessagesWithApiFallback( + chatId: string | number, + limit = 20, + offset = 0, +): Promise { + // 1. Try cache + const cached = await getMessages(chatId, limit, offset); + if (cached && cached.length > 0) return cached; + + // 2. Resolve chat for API call + const chat = getChatById(chatId); + if (!chat) return undefined; + + // 3. Rate limit before API call + try { + await enforceRateLimit("__api_fallback_messages"); + } catch { + return undefined; // Rate limited — return nothing rather than error + } + + // 4. Fetch from Telegram API + try { + const client = mtprotoService.getClient(); + const entity = chat.username ? chat.username : chat.id; + const inputPeer = await client.getInputEntity(entity); + + const result = await mtprotoService.withFloodWaitHandling(async () => { + return client.invoke( + new Api.messages.GetHistory({ + peer: inputPeer, + offsetId: 0, + offsetDate: 0, + addOffset: offset, + limit, + maxId: 0, + minId: 0, + hash: bigInt(0), + }), + ); + }); + + if ("messages" in result && Array.isArray(result.messages)) { + const messages = (result.messages as unknown as Record[]) + .filter((m) => m.className !== "MessageEmpty") + .map((m) => apiMessageToTelegramMessage(m, chat.id)); + + // Enrich fromName from users list + if ("users" in result && Array.isArray(result.users)) { + const usersById = new Map(); + for (const u of result.users as Array<{ + id?: unknown; + firstName?: string; + lastName?: string; + }>) { + if (u.id != null) { + const name = + [u.firstName, u.lastName].filter(Boolean).join(" ") || "Unknown"; + usersById.set(String(u.id), name); + } + } + for (const msg of messages) { + if (msg.fromId && !msg.fromName) { + msg.fromName = usersById.get(msg.fromId); + } + } + } + + return messages.length > 0 ? messages : undefined; + } + } catch { + // API failed — return nothing + } + + return undefined; +} + +/** Convert a raw GramJS dialog + chat/user to our TelegramChat format */ +function apiDialogToTelegramChat( + dialog: Record, + chatsById: Map>, + usersById: Map>, +): TelegramChat | undefined { + const peer = dialog.peer as { className?: string; userId?: unknown; chatId?: unknown; channelId?: unknown } | undefined; + if (!peer) return undefined; + + let id: string; + let type: TelegramChat["type"]; + let raw: Record | undefined; + + if (peer.className === "PeerUser" && peer.userId != null) { + id = String(peer.userId); + type = "private"; + raw = usersById.get(id); + } else if (peer.className === "PeerChat" && peer.chatId != null) { + id = String(peer.chatId); + type = "group"; + raw = chatsById.get(id); + } else if (peer.className === "PeerChannel" && peer.channelId != null) { + id = String(peer.channelId); + raw = chatsById.get(id); + type = raw && Boolean(raw.megagroup) ? "supergroup" : "channel"; + } else { + return undefined; + } + + let title: string; + let username: string | undefined; + if (raw) { + title = + (raw.title as string) ?? + [raw.firstName, raw.lastName].filter(Boolean).join(" ") ?? + "Unknown"; + username = raw.username as string | undefined; + } else { + title = "Unknown"; + } + + return { + id, + title, + type, + username, + unreadCount: typeof dialog.unreadCount === "number" ? dialog.unreadCount : 0, + isPinned: Boolean(dialog.pinned), + }; +} + +/** + * Get chats — tries cached Redux state first, falls back to Telegram API. + * Calls enforceRateLimit() internally before any API call. + */ +export async function getChatsWithApiFallback( + limit = 20, +): Promise { + // 1. Try cache + const cached = await getChats(limit); + if (cached.length > 0) return cached; + + // 2. Rate limit before API call + try { + await enforceRateLimit("__api_fallback_chats"); + } catch { + return []; + } + + // 3. Fetch from Telegram API + try { + const client = mtprotoService.getClient(); + + const result = await mtprotoService.withFloodWaitHandling(async () => { + return client.invoke( + new Api.messages.GetDialogs({ + offsetDate: 0, + offsetId: 0, + offsetPeer: new Api.InputPeerEmpty(), + limit, + hash: bigInt(0), + }), + ); + }); + + if (!("dialogs" in result) || !Array.isArray(result.dialogs)) return []; + + // Index chats and users by ID + const chatsById = new Map>(); + const usersById = new Map>(); + if ("chats" in result && Array.isArray(result.chats)) { + for (const c of result.chats as unknown as Record[]) { + if (c.id != null) chatsById.set(String(c.id), c); + } + } + if ("users" in result && Array.isArray(result.users)) { + for (const u of result.users as unknown as Record[]) { + if (u.id != null) usersById.set(String(u.id), u); + } + } + + return (result.dialogs as unknown as Record[]) + .map((d) => apiDialogToTelegramChat(d, chatsById, usersById)) + .filter((c): c is TelegramChat => c !== undefined) + .slice(0, limit); + } catch { + return []; + } +} + +/** + * Get current user — tries cached Redux state first, falls back to client.getMe(). + * Calls enforceRateLimit() internally before any API call. + */ +export async function getCurrentUserWithApiFallback(): Promise< + TelegramUser | undefined +> { + // 1. Try cache + const cached = getCurrentUser(); + if (cached) return cached; + + // 2. Rate limit before API call + try { + await enforceRateLimit("__api_fallback_me"); + } catch { + return undefined; + } + + // 3. Fetch from Telegram API + try { + const client = mtprotoService.getClient(); + const me = await mtprotoService.withFloodWaitHandling(async () => { + return client.getMe(); + }); + + if (!me) return undefined; + + const raw = me as unknown as { + id?: unknown; + firstName?: string; + lastName?: string; + username?: string; + phone?: string; + bot?: boolean; + }; + + return { + id: String(raw.id ?? ""), + firstName: raw.firstName ?? "", + lastName: raw.lastName, + username: raw.username, + phoneNumber: raw.phone, + isBot: Boolean(raw.bot), + }; + } catch { + return undefined; + } +} diff --git a/src/lib/mcp/telegram/tools/getChat.ts b/src/lib/mcp/telegram/tools/getChat.ts index f48d63aa6..52bd2e736 100644 --- a/src/lib/mcp/telegram/tools/getChat.ts +++ b/src/lib/mcp/telegram/tools/getChat.ts @@ -1,5 +1,7 @@ /** * Get Chat tool - Get detailed information about a specific chat + * + * Tries cached Redux state first, falls back to Telegram API when cache misses. */ import type { MCPTool, MCPToolResult } from "../../types"; @@ -8,6 +10,8 @@ import type { TelegramMCPContext } from "../types"; import { ErrorCategory, logAndFormatError } from "../../errorHandler"; import { formatEntity, getChatById } from "../telegramApi"; import { validateId } from "../../validation"; +import { mtprotoService } from "../../../../services/mtprotoService"; +import { enforceRateLimit } from "../../rateLimiter"; export const tool: MCPTool = { name: "get_chat", @@ -32,38 +36,65 @@ export async function getChat( const chatId = validateId(args.chat_id, "chat_id"); const chat = getChatById(chatId); - if (!chat) { + if (chat) { + // Cache hit — return cached data + return formatChatResult(chat); + } + + // Cache miss — try Telegram API + try { + await enforceRateLimit("__api_fallback_chat"); + + const client = mtprotoService.getClient(); + const entity = await mtprotoService.withFloodWaitHandling(async () => { + return client.getEntity(chatId); + }); + + if (!entity) { + return { + content: [{ type: "text", text: `Chat not found: ${chatId}` }], + isError: true, + }; + } + + const raw = entity as unknown as Record; + const result: string[] = []; + + result.push(`ID: ${raw.id ?? chatId}`); + + // Determine type and title + const className = raw.className as string | undefined; + if (className === "User") { + const name = + [raw.firstName, raw.lastName].filter(Boolean).join(" ") || "Unknown"; + result.push(`Name: ${name}`); + result.push(`Type: user`); + if (raw.username) result.push(`Username: @${raw.username}`); + if (raw.phone) result.push(`Phone: +${raw.phone}`); + if (raw.bot) result.push(`Bot: true`); + } else if (className === "Channel") { + result.push(`Title: ${raw.title ?? "Unknown"}`); + result.push(`Type: ${raw.megagroup ? "supergroup" : "channel"}`); + if (raw.username) result.push(`Username: @${raw.username}`); + if (raw.participantsCount) + result.push(`Participants: ${raw.participantsCount}`); + } else if (className === "Chat") { + result.push(`Title: ${raw.title ?? "Unknown"}`); + result.push(`Type: group`); + if (raw.participantsCount) + result.push(`Participants: ${raw.participantsCount}`); + } else { + result.push(`Title: ${raw.title ?? raw.firstName ?? "Unknown"}`); + result.push(`Type: unknown`); + } + + return { content: [{ type: "text", text: result.join("\n") }] }; + } catch { return { content: [{ type: "text", text: `Chat not found: ${chatId}` }], isError: true, }; } - - const entity = formatEntity(chat); - const result: string[] = []; - - result.push(`ID: ${entity.id}`); - result.push(`Title: ${entity.name}`); - result.push(`Type: ${entity.type}`); - if (entity.username) result.push(`Username: @${entity.username}`); - if ("participantsCount" in chat && chat.participantsCount) { - result.push(`Participants: ${chat.participantsCount}`); - } - if ("unreadCount" in chat) { - result.push(`Unread Messages: ${chat.unreadCount ?? 0}`); - } - - const lastMsg = chat.lastMessage; - if (lastMsg) { - const from = lastMsg.fromName ?? lastMsg.fromId ?? "Unknown"; - const date = new Date(lastMsg.date * 1000).toISOString(); - result.push(`Last Message: From ${from} at ${date}`); - result.push(`Message: ${lastMsg.message || "[Media/No text]"}`); - } - - return { - content: [{ type: "text", text: result.join("\n") }], - }; } catch (error) { return logAndFormatError( "get_chat", @@ -72,3 +103,31 @@ export async function getChat( ); } } + +function formatChatResult( + chat: ReturnType & object, +): MCPToolResult { + const entity = formatEntity(chat); + const result: string[] = []; + + result.push(`ID: ${entity.id}`); + result.push(`Title: ${entity.name}`); + result.push(`Type: ${entity.type}`); + if (entity.username) result.push(`Username: @${entity.username}`); + if ("participantsCount" in chat && chat.participantsCount) { + result.push(`Participants: ${chat.participantsCount}`); + } + if ("unreadCount" in chat) { + result.push(`Unread Messages: ${chat.unreadCount ?? 0}`); + } + + const lastMsg = (chat as { lastMessage?: { fromName?: string; fromId?: string; date: number; message?: string } }).lastMessage; + if (lastMsg) { + const from = lastMsg.fromName ?? lastMsg.fromId ?? "Unknown"; + const date = new Date(lastMsg.date * 1000).toISOString(); + result.push(`Last Message: From ${from} at ${date}`); + result.push(`Message: ${lastMsg.message || "[Media/No text]"}`); + } + + return { content: [{ type: "text", text: result.join("\n") }] }; +} diff --git a/src/lib/mcp/telegram/tools/getChats.ts b/src/lib/mcp/telegram/tools/getChats.ts index c0059c372..c6dd5f0ae 100644 --- a/src/lib/mcp/telegram/tools/getChats.ts +++ b/src/lib/mcp/telegram/tools/getChats.ts @@ -7,7 +7,7 @@ import type { TelegramMCPContext } from "../types"; import { ErrorCategory, logAndFormatError } from "../../errorHandler"; import { optNumber } from "../args"; -import { formatEntity, getChats as getChatsApi } from "../telegramApi"; +import { formatEntity, getChatsWithApiFallback } from "../telegramApi"; export const tool: MCPTool = { name: "get_chats", @@ -38,7 +38,7 @@ export async function getChats( const pageSize = optNumber(args, "page_size", 20); const start = (page - 1) * pageSize; - const chats = await getChatsApi(pageSize + start); + const chats = await getChatsWithApiFallback(pageSize + start); const paginatedChats = chats.slice(start, start + pageSize); if (paginatedChats.length === 0) { diff --git a/src/lib/mcp/telegram/tools/getHistory.ts b/src/lib/mcp/telegram/tools/getHistory.ts index 335441e0f..d1b811578 100644 --- a/src/lib/mcp/telegram/tools/getHistory.ts +++ b/src/lib/mcp/telegram/tools/getHistory.ts @@ -1,7 +1,7 @@ import type { MCPTool, MCPToolResult } from "../../types"; import type { TelegramMCPContext } from "../types"; import { ErrorCategory, logAndFormatError } from "../../errorHandler"; -import { getChatById, getMessages, formatMessage } from "../telegramApi"; +import { getChatById, getMessagesWithApiFallback, formatMessage } from "../telegramApi"; import { validateId } from "../../validation"; import { optNumber } from "../args"; @@ -34,7 +34,7 @@ export async function getHistory( }; } - const messages = await getMessages(chatId, limit, 0); + const messages = await getMessagesWithApiFallback(chatId, limit, 0); if (!messages || messages.length === 0) { return { content: [{ type: "text", text: "No messages found in this chat." }], diff --git a/src/lib/mcp/telegram/tools/getMe.ts b/src/lib/mcp/telegram/tools/getMe.ts index 36d588acf..0514fcba0 100644 --- a/src/lib/mcp/telegram/tools/getMe.ts +++ b/src/lib/mcp/telegram/tools/getMe.ts @@ -6,7 +6,7 @@ import type { MCPTool, MCPToolResult } from "../../types"; import type { TelegramMCPContext } from "../types"; import { ErrorCategory, logAndFormatError } from "../../errorHandler"; -import { formatEntity, getCurrentUser } from "../telegramApi"; +import { formatEntity, getCurrentUserWithApiFallback } from "../telegramApi"; export const tool: MCPTool = { name: "get_me", @@ -22,7 +22,7 @@ export async function getMe( _context: TelegramMCPContext, ): Promise { try { - const user = getCurrentUser(); + const user = await getCurrentUserWithApiFallback(); if (!user) { return { content: [{ type: "text", text: "User information not available" }], diff --git a/src/lib/mcp/telegram/tools/getMessageContext.ts b/src/lib/mcp/telegram/tools/getMessageContext.ts index 111db3fcb..eeccb7c49 100644 --- a/src/lib/mcp/telegram/tools/getMessageContext.ts +++ b/src/lib/mcp/telegram/tools/getMessageContext.ts @@ -1,7 +1,7 @@ import type { MCPTool, MCPToolResult } from "../../types"; import type { TelegramMCPContext } from "../types"; import { ErrorCategory, logAndFormatError } from "../../errorHandler"; -import { getChatById, getMessages, formatMessage } from "../telegramApi"; +import { getChatById, getMessagesWithApiFallback, formatMessage } from "../telegramApi"; import { validateId } from "../../validation"; import { optNumber } from "../args"; @@ -52,7 +52,7 @@ export async function getMessageContext( }; } - const allMessages = await getMessages(chatId, 200, 0); + const allMessages = await getMessagesWithApiFallback(chatId, 200, 0); if (!allMessages || allMessages.length === 0) { return { content: [{ type: "text", text: "No messages found in this chat." }], diff --git a/src/lib/mcp/telegram/tools/getMessages.ts b/src/lib/mcp/telegram/tools/getMessages.ts index c65a8a0b0..1c3fece64 100644 --- a/src/lib/mcp/telegram/tools/getMessages.ts +++ b/src/lib/mcp/telegram/tools/getMessages.ts @@ -9,7 +9,7 @@ import { ErrorCategory, logAndFormatError } from "../../errorHandler"; import { formatMessage, getChatById, - getMessages as getMessagesApi, + getMessagesWithApiFallback, } from "../telegramApi"; import { validateId } from "../../validation"; import type { TelegramMessage } from "../../../../store/telegram/types"; @@ -62,7 +62,7 @@ export async function getMessages( }; } - const messages = await getMessagesApi(chatId, pageSize, offset); + const messages = await getMessagesWithApiFallback(chatId, pageSize, offset); if (!messages || messages.length === 0) { return { content: [{ type: "text", text: "No messages found for this page." }], diff --git a/src/lib/mcp/telegram/tools/listChats.ts b/src/lib/mcp/telegram/tools/listChats.ts index 13c5fc64b..526bc26be 100644 --- a/src/lib/mcp/telegram/tools/listChats.ts +++ b/src/lib/mcp/telegram/tools/listChats.ts @@ -7,7 +7,7 @@ import type { TelegramMCPContext } from "../types"; import { ErrorCategory, logAndFormatError } from "../../errorHandler"; import { optNumber, optString } from "../args"; -import { formatEntity, getChats as getChatsApi } from "../telegramApi"; +import { formatEntity, getChatsWithApiFallback } from "../telegramApi"; import { toHumanReadableAction } from "../toolActionParser"; export const tool: MCPTool = { @@ -40,7 +40,7 @@ export async function listChats( const limit = optNumber(args, "limit", 20); const chatType = optString(args, "chat_type")?.toLowerCase(); - const chats = await getChatsApi(limit); + const chats = await getChatsWithApiFallback(limit); const contentItems: Array<{ type: "text"; text: string }> = []; for (const chat of chats) { diff --git a/src/lib/mcp/telegram/tools/listMessages.ts b/src/lib/mcp/telegram/tools/listMessages.ts index 03f6f36f7..fee0e9466 100644 --- a/src/lib/mcp/telegram/tools/listMessages.ts +++ b/src/lib/mcp/telegram/tools/listMessages.ts @@ -10,7 +10,7 @@ import { optNumber, optString } from "../args"; import { formatMessage, getChatById, - getMessages as getMessagesApi, + getMessagesWithApiFallback, } from "../telegramApi"; import { validateId } from "../../validation"; @@ -65,7 +65,7 @@ export async function listMessages( }; } - let messages = await getMessagesApi(chatId, limit * 2, 0); + let messages = await getMessagesWithApiFallback(chatId, limit * 2, 0); if (!messages || messages.length === 0) { return { content: [ diff --git a/src/lib/mcp/telegram/tools/resolveUsername.ts b/src/lib/mcp/telegram/tools/resolveUsername.ts index d8c11bd99..61c99db0f 100644 --- a/src/lib/mcp/telegram/tools/resolveUsername.ts +++ b/src/lib/mcp/telegram/tools/resolveUsername.ts @@ -1,5 +1,8 @@ /** * Resolve Username tool - Resolve a username to a user or chat ID + * + * Uses Telegram's contacts.ResolveUsername API for server-side resolution. + * Falls back to cached chat lookup if the API call fails. */ import type { MCPTool, MCPToolResult } from "../../types"; @@ -7,6 +10,8 @@ import type { TelegramMCPContext } from "../types"; import { ErrorCategory, logAndFormatError } from "../../errorHandler"; import { formatEntity, getChatById } from "../telegramApi"; +import { mtprotoService } from "../../../../services/mtprotoService"; +import { Api } from "telegram"; export const tool: MCPTool = { name: "resolve_username", @@ -35,11 +40,84 @@ export async function resolveUsername( isError: true, }; } - const username = raw.startsWith("@") ? raw : `@${raw}`; - const chat = getChatById(username); + const username = raw.startsWith("@") ? raw.slice(1) : raw; + + // Try server-side resolution via Telegram API + try { + const client = mtprotoService.getClient(); + + const result = await mtprotoService.withFloodWaitHandling(async () => { + return client.invoke( + new Api.contacts.ResolveUsername({ username }), + ); + }); + + if (result && "peer" in result) { + const peer = result.peer; + const peerType = + peer.className === "PeerUser" + ? "user" + : peer.className === "PeerChannel" + ? "channel" + : peer.className === "PeerChat" + ? "chat" + : "unknown"; + + // Extract ID from peer + const peerId = + "userId" in peer + ? String(peer.userId) + : "channelId" in peer + ? String(peer.channelId) + : "chatId" in peer + ? String(peer.chatId) + : "unknown"; + + // Try to get name from resolved users/chats + let name = username; + if ("users" in result && Array.isArray(result.users)) { + const user = result.users[0] as { + firstName?: string; + lastName?: string; + }; + if (user) { + name = + [user.firstName, user.lastName].filter(Boolean).join(" ") || + username; + } + } + if ("chats" in result && Array.isArray(result.chats)) { + const chat = result.chats[0] as { title?: string }; + if (chat?.title) { + name = chat.title; + } + } + + return { + content: [ + { + type: "text", + text: JSON.stringify( + { id: peerId, name, type: peerType, username }, + undefined, + 2, + ), + }, + ], + }; + } + } catch { + // API call failed — fall back to cached lookup below + } + + // Fallback: look up from cached state + const lookupKey = `@${username}`; + const chat = getChatById(lookupKey); if (!chat) { return { - content: [{ type: "text", text: `Username ${username} not found` }], + content: [ + { type: "text", text: `Username @${username} not found` }, + ], isError: true, }; } diff --git a/src/lib/mcp/telegram/tools/searchMessages.ts b/src/lib/mcp/telegram/tools/searchMessages.ts index 627b763a4..8f1a4c46c 100644 --- a/src/lib/mcp/telegram/tools/searchMessages.ts +++ b/src/lib/mcp/telegram/tools/searchMessages.ts @@ -1,5 +1,8 @@ /** * Search Messages tool - Search for messages in a chat by text + * + * Uses Telegram's messages.Search API for server-side full-text search. + * Falls back to filtering cached messages if the API call fails. */ import type { MCPTool, MCPToolResult } from "../../types"; @@ -9,6 +12,11 @@ import { ErrorCategory, logAndFormatError } from "../../errorHandler"; import { optNumber } from "../args"; import { formatMessage, getChatById, getMessages } from "../telegramApi"; import { validateId } from "../../validation"; +import { mtprotoService } from "../../../../services/mtprotoService"; +import { Api } from "telegram"; +import bigInt from "big-integer"; +import type { ApiMessage } from "../apiResultTypes"; +import { narrow } from "../apiCastHelpers"; export const tool: MCPTool = { name: "search_messages", @@ -52,6 +60,57 @@ export async function searchMessages( }; } + // Try server-side search via Telegram API + try { + const client = mtprotoService.getClient(); + const entity = chat.username ? chat.username : chat.id; + const inputPeer = await client.getInputEntity(entity); + + const result = await mtprotoService.withFloodWaitHandling(async () => { + return client.invoke( + new Api.messages.Search({ + peer: inputPeer, + q: query, + filter: new Api.InputMessagesFilterEmpty(), + minDate: 0, + maxDate: 0, + offsetId: 0, + addOffset: 0, + limit, + maxId: 0, + minId: 0, + hash: bigInt(0), + }), + ); + }); + + if ("messages" in result && Array.isArray(result.messages)) { + const messages = narrow(result.messages); + + if (messages.length === 0) { + return { + content: [ + { type: "text", text: `No messages matching "${query}" found.` }, + ], + }; + } + + const lines = messages.map((msg) => { + const id = msg.id ?? "?"; + const text = msg.message ?? "[Media/No text]"; + const date = msg.date + ? new Date(msg.date * 1000).toISOString() + : "unknown"; + return `ID: ${id} | Date: ${date} | ${text}`; + }); + + return { content: [{ type: "text", text: lines.join("\n") }] }; + } + } catch { + // API call failed — fall back to cached message search below + } + + // Fallback: search cached messages locally const messages = await getMessages(chatId, Math.min(limit * 3, 100), 0); if (!messages || messages.length === 0) { return { content: [{ type: "text", text: "No messages found." }] }; @@ -62,12 +121,24 @@ export async function searchMessages( .filter((m) => (m.message ?? "").toLowerCase().includes(q)) .slice(0, limit); + if (filtered.length === 0) { + return { + content: [ + { type: "text", text: `No messages matching "${query}" found.` }, + ], + }; + } + const lines = filtered.map((m) => { const f = formatMessage(m); return `ID: ${f.id} | ${m.fromName ?? m.fromId ?? "Unknown"} | ${f.date} | ${f.text || "[Media]"}`; }); - return { content: [{ type: "text", text: lines.join("\n") }] }; + return { + content: [ + { type: "text", text: `(cached search)\n${lines.join("\n")}` }, + ], + }; } catch (error) { return logAndFormatError( "search_messages", diff --git a/src/lib/mcp/telegram/tools/searchPublicChats.ts b/src/lib/mcp/telegram/tools/searchPublicChats.ts index da6a5036d..4eec984f1 100644 --- a/src/lib/mcp/telegram/tools/searchPublicChats.ts +++ b/src/lib/mcp/telegram/tools/searchPublicChats.ts @@ -1,5 +1,8 @@ /** * Search Public Chats tool - Search for public chats, channels, or bots + * + * Uses Telegram's contacts.Search API for server-side discovery. + * Falls back to filtering cached chats if the API call fails. */ import type { MCPTool, MCPToolResult } from "../../types"; @@ -7,6 +10,8 @@ import type { TelegramMCPContext } from "../types"; import { ErrorCategory, logAndFormatError } from "../../errorHandler"; import { formatEntity, searchChats } from "../telegramApi"; +import { mtprotoService } from "../../../../services/mtprotoService"; +import { Api } from "telegram"; export const tool: MCPTool = { name: "search_public_chats", @@ -33,10 +38,102 @@ export async function searchPublicChats( isError: true, }; } + + // Try server-side search via Telegram API + try { + const client = mtprotoService.getClient(); + + const result = await mtprotoService.withFloodWaitHandling(async () => { + return client.invoke( + new Api.contacts.Search({ q: query, limit: 20 }), + ); + }); + + const entries: Array<{ + id: string; + name: string; + type: string; + username?: string; + }> = []; + + // Process returned chats + if ("chats" in result && Array.isArray(result.chats)) { + for (const chat of result.chats) { + const c = chat as { + id: { toString(): string }; + title?: string; + username?: string; + megagroup?: boolean; + broadcast?: boolean; + }; + const type = c.broadcast + ? "channel" + : c.megagroup + ? "group" + : "chat"; + entries.push({ + id: String(c.id), + name: c.title ?? "Unknown", + type, + username: c.username, + }); + } + } + + // Process returned users + if ("users" in result && Array.isArray(result.users)) { + for (const user of result.users) { + const u = user as { + id: { toString(): string }; + firstName?: string; + lastName?: string; + username?: string; + bot?: boolean; + }; + const name = + [u.firstName, u.lastName].filter(Boolean).join(" ") || "Unknown"; + entries.push({ + id: String(u.id), + name, + type: u.bot ? "bot" : "user", + username: u.username, + }); + } + } + + if (entries.length > 0) { + return { + content: [ + { type: "text", text: JSON.stringify(entries, undefined, 2) }, + ], + }; + } + } catch { + // API call failed — fall back to cached search below + } + + // Fallback: search cached chats locally const chats = await searchChats(query); const results = chats.map(formatEntity); + + if (results.length === 0) { + return { + content: [ + { + type: "text", + text: `No public chats matching "${query}" found.`, + }, + ], + }; + } + return { - content: [{ type: "text", text: JSON.stringify(results, undefined, 2) }], + content: [ + { + type: "text", + text: `(cached search)\n${JSON.stringify(results, undefined, 2)}`, + }, + ], }; } catch (error) { return logAndFormatError(