diff --git a/main.py b/main.py index 9721118..00f7db7 100644 --- a/main.py +++ b/main.py @@ -310,6 +310,29 @@ def format_entity(entity) -> Dict[str, Any]: return result +async def resolve_entity(identifier: Union[int, str]) -> Any: + """Resolve entity with automatic cache warming on miss. + + StringSession has no persistent entity cache. If get_entity() fails + because the cache is cold (ValueError on PeerUser lookup for group IDs), + warm the cache via get_dialogs() and retry. + """ + try: + return await client.get_entity(identifier) + except ValueError: + await client.get_dialogs() + return await client.get_entity(identifier) + + +async def resolve_input_entity(identifier: Union[int, str]) -> Any: + """Like resolve_entity() but returns an InputPeer.""" + try: + return await client.get_input_entity(identifier) + except ValueError: + await client.get_dialogs() + return await client.get_input_entity(identifier) + + def format_message(message) -> Dict[str, Any]: """Helper function to format message information consistently.""" result = { @@ -400,7 +423,7 @@ async def get_messages(chat_id: Union[int, str], page: int = 1, page_size: int = page_size: Number of messages per page. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) offset = (page - 1) * page_size messages = await client.get_messages(entity, limit=page_size, add_offset=offset) if not messages: @@ -436,7 +459,7 @@ async def send_message(chat_id: Union[int, str], message: str) -> str: message: The message content to send. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.send_message(entity, message) return "Message sent successfully." except Exception as e: @@ -457,7 +480,7 @@ async def subscribe_public_channel(channel: Union[int, str]) -> str: Subscribe (join) to a public channel or supergroup by username or ID. """ try: - entity = await client.get_entity(channel) + entity = await resolve_entity(channel) await client(functions.channels.JoinChannelRequest(channel=entity)) title = getattr(entity, "title", getattr(entity, "username", "Unknown channel")) return f"Subscribed to {title}." @@ -487,7 +510,7 @@ async def list_inline_buttons( else: return "message_id must be an integer." - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) target_message = None if message_id is not None: @@ -574,7 +597,7 @@ async def press_inline_button( else: return "button_index must be an integer." - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) target_message = None if message_id is not None: @@ -751,7 +774,7 @@ async def list_messages( to_date: Filter messages until this date (format: YYYY-MM-DD). """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) # Parse date filters if provided from_date_obj = None @@ -876,7 +899,7 @@ async def list_topics( search_query: Optional query to filter topics by title. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) if not isinstance(entity, Channel) or not getattr(entity, "megagroup", False): return "The specified chat is not a supergroup." @@ -1016,7 +1039,7 @@ async def get_chat(chat_id: Union[int, str]) -> str: chat_id: The ID or username of the chat. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) result = [] result.append(f"ID: {entity.id}") @@ -1150,7 +1173,7 @@ async def get_contact_chats(contact_id: Union[int, str]) -> str: """ try: # Get contact info - contact = await client.get_entity(contact_id) + contact = await resolve_entity(contact_id) if not isinstance(contact, User): return f"ID {contact_id} is not a user/contact." @@ -1207,7 +1230,7 @@ async def get_last_interaction(contact_id: Union[int, str]) -> str: """ try: # Get contact info - contact = await client.get_entity(contact_id) + contact = await resolve_entity(contact_id) if not isinstance(contact, User): return f"ID {contact_id} is not a user/contact." @@ -1249,7 +1272,7 @@ async def get_message_context( context_size: Number of messages before and after to include. """ try: - chat = await client.get_entity(chat_id) + chat = await resolve_entity(chat_id) # Get messages around the specified message messages_before = await client.get_messages(chat, limit=context_size, max_id=message_id) central_message = await client.get_messages(chat, ids=message_id) @@ -1448,7 +1471,7 @@ async def delete_contact(user_id: Union[int, str]) -> str: user_id: The Telegram user ID or username of the contact to delete. """ try: - user = await client.get_entity(user_id) + user = await resolve_entity(user_id) await client(functions.contacts.DeleteContactsRequest(id=[user])) return f"Contact with user ID {user_id} deleted." except Exception as e: @@ -1468,7 +1491,7 @@ async def block_user(user_id: Union[int, str]) -> str: user_id: The Telegram user ID or username to block. """ try: - user = await client.get_entity(user_id) + user = await resolve_entity(user_id) await client(functions.contacts.BlockRequest(id=user)) return f"User {user_id} blocked." except Exception as e: @@ -1488,7 +1511,7 @@ async def unblock_user(user_id: Union[int, str]) -> str: user_id: The Telegram user ID or username to unblock. """ try: - user = await client.get_entity(user_id) + user = await resolve_entity(user_id) await client(functions.contacts.UnblockRequest(id=user)) return f"User {user_id} unblocked." except Exception as e: @@ -1524,7 +1547,7 @@ async def create_group(title: str, user_ids: List[Union[int, str]]) -> str: users = [] for user_id in user_ids: try: - user = await client.get_entity(user_id) + user = await resolve_entity(user_id) users.append(user) except Exception as e: logger.error(f"Failed to get entity for user ID {user_id}: {e}") @@ -1583,12 +1606,12 @@ async def invite_to_group(group_id: Union[int, str], user_ids: List[Union[int, s user_ids: List of user IDs or usernames to invite. """ try: - entity = await client.get_entity(group_id) + entity = await resolve_entity(group_id) users_to_add = [] for user_id in user_ids: try: - user = await client.get_entity(user_id) + user = await resolve_entity(user_id) users_to_add.append(user) except ValueError as e: return f"Error: User with ID {user_id} could not be found. {e}" @@ -1636,7 +1659,7 @@ async def leave_chat(chat_id: Union[int, str]) -> str: chat_id: The chat ID or username to leave. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) # Check the entity type carefully if isinstance(entity, Channel): @@ -1743,7 +1766,7 @@ async def send_file(chat_id: Union[int, str], file_path: str, caption: str = Non return f"File not found: {file_path}" if not os.access(file_path, os.R_OK): return f"File is not readable: {file_path}" - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.send_file(entity, file_path, caption=caption) return f"File sent to chat {chat_id}." except Exception as e: @@ -1765,7 +1788,7 @@ async def download_media(chat_id: Union[int, str], message_id: int, file_path: s file_path: Absolute path to save the downloaded file (must be writable). """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) msg = await client.get_messages(entity, ids=message_id) if not msg or not msg.media: return "No media found in the specified message." @@ -1933,7 +1956,7 @@ async def set_privacy_settings( allow_entities = [] for user_id in allow_users: try: - user = await client.get_entity(user_id) + user = await resolve_entity(user_id) allow_entities.append(user) except Exception as user_err: logger.warning(f"Could not get entity for user ID {user_id}: {user_err}") @@ -1950,7 +1973,7 @@ async def set_privacy_settings( disallow_entities = [] for user_id in disallow_users: try: - user = await client.get_entity(user_id) + user = await resolve_entity(user_id) disallow_entities.append(user) except Exception as user_err: logger.warning(f"Could not get entity for user ID {user_id}: {user_err}") @@ -2058,7 +2081,7 @@ async def edit_chat_title(chat_id: Union[int, str], title: str) -> str: Edit the title of a chat, group, or channel. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) if isinstance(entity, Channel): await client(functions.channels.EditTitleRequest(channel=entity, title=title)) elif isinstance(entity, Chat): @@ -2087,7 +2110,7 @@ async def edit_chat_photo(chat_id: Union[int, str], file_path: str) -> str: if not os.access(file_path, os.R_OK): return f"Photo file not readable: {file_path}" - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) uploaded_file = await client.upload_file(file_path) if isinstance(entity, Channel): @@ -2120,7 +2143,7 @@ async def delete_chat_photo(chat_id: Union[int, str]) -> str: Delete the photo of a chat, group, or channel. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) if isinstance(entity, Channel): # Use InputChatPhotoEmpty for channels/supergroups await client( @@ -2160,8 +2183,8 @@ async def promote_admin( rights: Admin rights to give (optional) """ try: - chat = await client.get_entity(group_id) - user = await client.get_entity(user_id) + chat = await resolve_entity(group_id) + user = await resolve_entity(user_id) # Set default admin rights if not provided if not rights: @@ -2228,8 +2251,8 @@ async def demote_admin(group_id: Union[int, str], user_id: Union[int, str]) -> s user_id: User ID or username to demote """ try: - chat = await client.get_entity(group_id) - user = await client.get_entity(user_id) + chat = await resolve_entity(group_id) + user = await resolve_entity(user_id) # Create empty admin rights (regular user) admin_rights = ChatAdminRights( @@ -2281,8 +2304,8 @@ async def ban_user(chat_id: Union[int, str], user_id: Union[int, str]) -> str: user_id: User ID or username to ban """ try: - chat = await client.get_entity(chat_id) - user = await client.get_entity(user_id) + chat = await resolve_entity(chat_id) + user = await resolve_entity(user_id) # Create banned rights (all restrictions enabled) banned_rights = ChatBannedRights( @@ -2332,8 +2355,8 @@ async def unban_user(chat_id: Union[int, str], user_id: Union[int, str]) -> str: user_id: User ID or username to unban """ try: - chat = await client.get_entity(chat_id) - user = await client.get_entity(user_id) + chat = await resolve_entity(chat_id) + user = await resolve_entity(user_id) # Create unbanned rights (no restrictions) unbanned_rights = ChatBannedRights( @@ -2419,7 +2442,7 @@ async def get_invite_link(chat_id: Union[int, str]) -> str: Get the invite link for a group or channel. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) # Try using ExportChatInviteRequest first try: @@ -2513,7 +2536,7 @@ async def export_chat_invite(chat_id: Union[int, str]) -> str: Export a chat invite link. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) # Try using ExportChatInviteRequest first try: @@ -2631,7 +2654,7 @@ async def send_voice(chat_id: Union[int, str], file_path: str) -> str: ): return "Voice file must be .ogg or .opus format." - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.send_file(entity, file_path, voice_note=True) return f"Voice message sent to chat {chat_id}." except Exception as e: @@ -2649,8 +2672,8 @@ async def forward_message( Forward a message from one chat to another. """ try: - from_entity = await client.get_entity(from_chat_id) - to_entity = await client.get_entity(to_chat_id) + from_entity = await resolve_entity(from_chat_id) + to_entity = await resolve_entity(to_chat_id) await client.forward_messages(to_entity, message_id, from_entity) return f"Message {message_id} forwarded from {from_chat_id} to {to_chat_id}." except Exception as e: @@ -2674,7 +2697,7 @@ async def edit_message(chat_id: Union[int, str], message_id: int, new_text: str) Edit a message you sent. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.edit_message(entity, message_id, new_text) return f"Message {message_id} edited." except Exception as e: @@ -2694,7 +2717,7 @@ async def delete_message(chat_id: Union[int, str], message_id: int) -> str: Delete a message by ID. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.delete_messages(entity, message_id) return f"Message {message_id} deleted." except Exception as e: @@ -2712,7 +2735,7 @@ async def pin_message(chat_id: Union[int, str], message_id: int) -> str: Pin a message in a chat. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.pin_message(entity, message_id) return f"Message {message_id} pinned in chat {chat_id}." except Exception as e: @@ -2730,7 +2753,7 @@ async def unpin_message(chat_id: Union[int, str], message_id: int) -> str: Unpin a message in a chat. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.unpin_message(entity, message_id) return f"Message {message_id} unpinned in chat {chat_id}." except Exception as e: @@ -2748,7 +2771,7 @@ async def mark_as_read(chat_id: Union[int, str]) -> str: Mark all messages as read in a chat. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.send_read_acknowledge(entity) return f"Marked all messages as read in chat {chat_id}." except Exception as e: @@ -2764,7 +2787,7 @@ async def reply_to_message(chat_id: Union[int, str], message_id: int, text: str) Reply to a specific message in a chat. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.send_message(entity, text, reply_to=message_id) return f"Replied to message {message_id} in chat {chat_id}." except Exception as e: @@ -2786,7 +2809,7 @@ async def get_media_info(chat_id: Union[int, str], message_id: int) -> str: message_id: The message ID. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) msg = await client.get_messages(entity, ids=message_id) if not msg or not msg.media: @@ -2820,7 +2843,7 @@ async def search_messages(chat_id: Union[int, str], query: str, limit: int = 20) Search for messages in a chat by text. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) messages = await client.get_messages(entity, limit=limit, search=query) lines = [] @@ -2866,7 +2889,7 @@ async def mute_chat(chat_id: Union[int, str]) -> str: try: from telethon.tl.types import InputPeerNotifySettings - peer = await client.get_entity(chat_id) + peer = await resolve_entity(chat_id) await client( functions.account.UpdateNotifySettingsRequest( peer=peer, settings=InputPeerNotifySettings(mute_until=2**31 - 1) @@ -2876,7 +2899,7 @@ async def mute_chat(chat_id: Union[int, str]) -> str: except (ImportError, AttributeError) as type_err: try: # Alternative approach directly using raw API - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) await client( functions.account.UpdateNotifySettingsRequest( peer=peer, @@ -2909,7 +2932,7 @@ async def unmute_chat(chat_id: Union[int, str]) -> str: try: from telethon.tl.types import InputPeerNotifySettings - peer = await client.get_entity(chat_id) + peer = await resolve_entity(chat_id) await client( functions.account.UpdateNotifySettingsRequest( peer=peer, settings=InputPeerNotifySettings(mute_until=0) @@ -2919,7 +2942,7 @@ async def unmute_chat(chat_id: Union[int, str]) -> str: except (ImportError, AttributeError) as type_err: try: # Alternative approach directly using raw API - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) await client( functions.account.UpdateNotifySettingsRequest( peer=peer, @@ -2950,7 +2973,7 @@ async def archive_chat(chat_id: Union[int, str]) -> str: Archive a chat. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) peer = utils.get_input_peer(entity) await client( functions.folders.EditPeerFoldersRequest( @@ -2973,7 +2996,7 @@ async def unarchive_chat(chat_id: Union[int, str]) -> str: Unarchive a chat. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) peer = utils.get_input_peer(entity) await client( functions.folders.EditPeerFoldersRequest( @@ -3019,7 +3042,7 @@ async def send_sticker(chat_id: Union[int, str], file_path: str) -> str: if not file_path.lower().endswith(".webp"): return "Sticker file must be a .webp file." - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.send_file(entity, file_path, force_document=False) return f"Sticker sent to chat {chat_id}." except Exception as e: @@ -3097,7 +3120,7 @@ async def send_gif(chat_id: Union[int, str], gif_id: int) -> str: try: if not isinstance(gif_id, int): return "gif_id must be a Telegram document ID (integer), not a file path. Use get_gif_search to find IDs." - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) await client.send_file(entity, gif_id) return f"GIF sent to chat {chat_id}." except Exception as e: @@ -3110,7 +3133,7 @@ async def get_bot_info(bot_username: str) -> str: Get information about a bot by username. """ try: - entity = await client.get_entity(bot_username) + entity = await resolve_entity(bot_username) if not entity: return f"Bot with username {bot_username} not found." @@ -3171,7 +3194,7 @@ async def set_bot_commands(bot_username: str, commands: list) -> str: ] # Get the bot entity - bot = await client.get_entity(bot_username) + bot = await resolve_entity(bot_username) # Set the commands with proper scope await client( @@ -3198,7 +3221,7 @@ async def get_history(chat_id: Union[int, str], limit: int = 100) -> str: Get full chat history (up to limit). """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) messages = await client.get_messages(entity, limit=limit) lines = [] @@ -3224,7 +3247,7 @@ async def get_user_photos(user_id: Union[int, str], limit: int = 10) -> str: Get profile photos of a user. """ try: - user = await client.get_entity(user_id) + user = await resolve_entity(user_id) photos = await client( functions.photos.GetUserPhotosRequest(user_id=user, offset=0, max_id=0, limit=limit) ) @@ -3242,7 +3265,7 @@ async def get_user_status(user_id: Union[int, str]) -> str: Get the online status of a user. """ try: - user = await client.get_entity(user_id) + user = await resolve_entity(user_id) return str(user.status) except Exception as e: return log_and_format_error("get_user_status", e, user_id=user_id) @@ -3288,7 +3311,7 @@ async def get_pinned_messages(chat_id: Union[int, str]) -> str: Get all pinned messages in a chat. """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) # Use correct filter based on Telethon version try: @@ -3345,7 +3368,7 @@ async def create_poll( close_date: Optional close date in ISO format (YYYY-MM-DD HH:MM:SS) """ try: - entity = await client.get_entity(chat_id) + entity = await resolve_entity(chat_id) # Validate options if len(options) < 2: @@ -3419,7 +3442,7 @@ async def send_reaction( try: from telethon.tl.types import ReactionEmoji - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) await client( functions.messages.SendReactionRequest( peer=peer, @@ -3456,7 +3479,7 @@ async def remove_reaction( message_id: The message ID to remove reaction from """ try: - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) await client( functions.messages.SendReactionRequest( peer=peer, @@ -3492,7 +3515,7 @@ async def get_message_reactions( try: from telethon.tl.types import ReactionEmoji, ReactionCustomEmoji - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) result = await client( functions.messages.GetMessageReactionsListRequest( @@ -3569,7 +3592,7 @@ async def save_draft( no_webpage: If True, disable link preview in the draft """ try: - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) # Build reply_to parameter if provided reply_to = None @@ -3664,7 +3687,7 @@ async def clear_draft(chat_id: Union[int, str]) -> str: chat_id: The chat ID or username to clear the draft from """ try: - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) # Saving an empty message clears the draft await client( @@ -3760,7 +3783,7 @@ async def get_folder(folder_id: int) -> str: included_chats = [] for peer in getattr(target_folder, "include_peers", []): try: - entity = await client.get_entity(peer) + entity = await resolve_entity(peer) chat_info = { "id": entity.id, "name": getattr(entity, "title", None) @@ -3777,7 +3800,7 @@ async def get_folder(folder_id: int) -> str: excluded_chats = [] for peer in getattr(target_folder, "exclude_peers", []): try: - entity = await client.get_entity(peer) + entity = await resolve_entity(peer) chat_info = { "id": entity.id, "name": getattr(entity, "title", None) @@ -3792,7 +3815,7 @@ async def get_folder(folder_id: int) -> str: pinned_chats = [] for peer in getattr(target_folder, "pinned_peers", []): try: - entity = await client.get_entity(peer) + entity = await resolve_entity(peer) chat_info = { "id": entity.id, "name": getattr(entity, "title", None) @@ -3892,7 +3915,7 @@ async def create_folder( if chat_ids: for chat_id in chat_ids: try: - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) include_peers.append(peer) except Exception as e: return f"Failed to resolve chat '{chat_id}': {str(e)}" @@ -3967,7 +3990,7 @@ async def add_chat_to_folder( # Resolve chat to input peer try: - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) except Exception as e: return f"Failed to resolve chat '{chat_id}': {str(e)}" @@ -4057,7 +4080,7 @@ async def remove_chat_from_folder(folder_id: int, chat_id: Union[int, str]) -> s # Resolve chat to get peer ID try: - peer = await client.get_input_entity(chat_id) + peer = await resolve_input_entity(chat_id) peer_id = utils.get_peer_id(peer) except Exception as e: return f"Failed to resolve chat '{chat_id}': {str(e)}" @@ -4214,6 +4237,11 @@ async def _main() -> None: print("Starting Telegram client...") await client.start() + # Warm entity cache — StringSession has no persistent cache, + # so fetch all dialogs once to populate it + print("Warming entity cache...") + await client.get_dialogs() + print("Telegram client started. Running MCP server...") # Use the asynchronous entrypoint instead of mcp.run() await mcp.run_stdio_async()