Merge pull request chigwell#173 from ex3lite/feat/contact-aliases-and-qr-fixes

feat(contacts): favorite aliases resolved everywhere; fix session generator 2FA retry and non-interactive stdin

Author: chigwell

.gitignore

@@ -199,3 +199,4 @@ test_upload.txt
 test_voice.ogg
 sticker.webp
 two.png
+aliases.json

README.md

@@ -46,7 +46,7 @@ The server currently includes 80+ MCP tools grouped into these areas:
 - **Accounts:** list configured accounts and route tool calls by account label.
 - **Chats and groups:** list chats, inspect metadata, create groups/channels, join or leave chats, invite users, manage admins, bans, default permissions, slow mode, topics, invite links, common chats, read receipts, and message links.
 - **Messages:** send, schedule, edit, delete, forward, pin, unpin, mark read, reply, search, inspect context, create polls, manage reactions, inspect inline buttons, and press inline callbacks.
-- **Contacts:** list, search, add, delete, block, unblock, import, export, inspect direct chats, and find recent contact interactions.
+- **Contacts:** list, search, add, delete, block, unblock, import, export, inspect direct chats, find recent contact interactions, and manage favorite aliases (e.g. save "andrew" so any tool accepting a `chat_id` resolves it; searches check favorites first).
 - **Media:** send files, download media, upload files, send voice notes, stickers, GIFs, and inspect message media.
 - **Profile and privacy:** get your own account info, update profile fields, set or delete profile photos, inspect privacy settings, get user info/photos/status, and manage bot commands.
 - **Folders and drafts:** list, create, update, reorder, and delete Telegram folders; save, list, and clear drafts.

session_string_generator.py

@@ -108,11 +108,15 @@ def _qr_login(client: TelegramClient) -> None:
             print("\nQR code expired, here is a fresh one.")
             _render_qr(qr)
         except errors.SessionPasswordNeededError:
-            pw = getpass.getpass(
-                "\nTwo-factor authentication enabled. Please enter your password: "
-            )
-            client.sign_in(password=pw)
-            return
+            while True:
+                pw = getpass.getpass(
+                    "\nTwo-factor authentication enabled. Please enter your password: "
+                )
+                try:
+                    client.sign_in(password=pw)
+                    return
+                except errors.PasswordHashInvalidError:
+                    print("Invalid password, please try again.")
 
     print("\nQR code expired too many times. Please run the generator again.")
     client.disconnect()
@@ -170,11 +174,15 @@ def main() -> None:
         "\nYour credentials will NOT be stored on any server and are only used for local authentication.\n"
     )
 
-    label = (
-        input("Account label (optional, e.g. 'work', 'personal'; leave empty for default): ")
-        .strip()
-        .lower()
-    )
+    try:
+        label = (
+            input("Account label (optional, e.g. 'work', 'personal'; leave empty for default): ")
+            .strip()
+            .lower()
+        )
+    except EOFError:
+        # Non-interactive stdin (piped/scripted runs): fall back to the default label.
+        label = ""
 
     if args.qr:
         method = "1"
@@ -210,9 +218,12 @@ def main() -> None:
         print(f"{env_var}={session_string}")
         print("\nIMPORTANT: Keep this string private and never share it with anyone!")
 
-        choice = input(
-            "\nWould you like to automatically update your .env file with this session string? (y/N): "
-        )
+        try:
+            choice = input(
+                "\nWould you like to automatically update your .env file with this session string? (y/N): "
+            )
+        except EOFError:
+            choice = "n"
         if choice.lower() == "y":
             try:
                 with open(".env", "r") as file:

telegram_mcp/runtime.py

@@ -816,6 +816,31 @@ def format_entity(entity) -> Dict[str, Any]:
     return result
 
 
+_ALIASES_FILE = Path(__file__).resolve().parent.parent / "aliases.json"
+
+
+def load_aliases() -> Dict[str, int]:
+    try:
+        with open(_ALIASES_FILE, "r", encoding="utf-8") as f:
+            return {k.lower(): int(v) for k, v in json.load(f).items()}
+    except (FileNotFoundError, ValueError):
+        return {}
+
+
+def save_aliases(aliases: Dict[str, int]) -> None:
+    with open(_ALIASES_FILE, "w", encoding="utf-8") as f:
+        json.dump(aliases, f, ensure_ascii=False, indent=2)
+
+
+def apply_alias(identifier: Union[int, str]) -> Union[int, str]:
+    """If identifier matches a saved alias (case-insensitive), return its chat ID."""
+    if isinstance(identifier, str):
+        alias_id = load_aliases().get(identifier.strip().lstrip("@").lower())
+        if alias_id is not None:
+            return alias_id
+    return identifier
+
+
 def _marked_id_candidates(identifier: Union[int, str]) -> list[int]:
     """Return marked chat/channel ID variants for a bare positive integer ID."""
     if not isinstance(identifier, int) or identifier <= 0:
@@ -839,6 +864,7 @@ async def resolve_entity(identifier: Union[int, str], client=None) -> Any:
 
     On ConnectionError, reconnects and retries once.
     """
+    identifier = apply_alias(identifier)
     if client is None:
         client = get_client()
     await ensure_connected(client)
@@ -882,6 +908,7 @@ async def resolve_input_entity(identifier: Union[int, str], client=None) -> Any:
 
     Uses the same cache warming, marked-ID fallback, and reconnect behavior.
     """
+    identifier = apply_alias(identifier)
     if client is None:
         client = get_client()
     await ensure_connected(client)

telegram_mcp/tools/contacts.py

@@ -47,6 +47,7 @@ async def list_contacts(account: Optional[str] = None) -> str:
 async def search_contacts(query: str, account: Optional[str] = None) -> str:
     """
     Search for contacts by name, username, or phone number using Telethon's SearchRequest.
+    Saved favorite aliases matching the query are checked first and returned at the top.
     Args:
         query: The search term to look for in contact names, usernames, or phone numbers.
 
@@ -55,11 +56,17 @@ async def search_contacts(query: str, account: Optional[str] = None) -> str:
     try:
         cl = get_client(account)
         await ensure_connected(cl)
+        q = query.strip().lstrip("@").lower()
+        alias_records = [
+            {"alias": alias, "id": chat_id, "favorite": True}
+            for alias, chat_id in load_aliases().items()
+            if q and (q in alias or alias in q)
+        ]
         result = await cl(functions.contacts.SearchRequest(q=query, limit=50))
         users = result.users
-        if not users:
+        if not users and not alias_records:
             return f"No contacts found matching '{query}'."
-        records = []
+        records = alias_records
         for user in users:
             name = f"{getattr(user, 'first_name', '')} {getattr(user, 'last_name', '')}".strip()
             record = {
@@ -596,6 +603,57 @@ async def send_contact(
         return log_and_format_error("send_contact", e, chat_id=chat_id, phone_number=phone_number)
 
 
+@mcp.tool(annotations=ToolAnnotations(title="Set Contact Alias", openWorldHint=True))
+@with_account(readonly=True)
+async def set_contact_alias(alias: str, chat_id: str, account: Optional[str] = None) -> str:
+    """
+    Save a favorite alias for a contact or chat. After this, the alias can be used
+    anywhere a chat_id is accepted (e.g. alias "андрей" -> send_message("андрей", ...)).
+    Args:
+        alias: Short name to remember (case-insensitive, e.g. "андрей").
+        chat_id: Chat ID, username (@user), or phone of the target.
+    """
+    try:
+        cl = get_client(account)
+        entity = await resolve_entity(chat_id, cl)
+        marked_id = get_marked_id(entity)
+        aliases = load_aliases()
+        aliases[alias.strip().lstrip("@").lower()] = marked_id
+        save_aliases(aliases)
+        return format_tool_result(
+            {"alias": alias.strip().lower(), "resolved": format_entity(entity)}
+        )
+    except Exception as e:
+        return log_and_format_error("set_contact_alias", e, alias=alias, chat_id=chat_id)
+
+
+@mcp.tool(annotations=ToolAnnotations(title="List Contact Aliases", readOnlyHint=True))
+@with_account(readonly=True)
+async def list_contact_aliases(account: Optional[str] = None) -> str:
+    """List all saved favorite aliases and their chat IDs."""
+    try:
+        aliases = load_aliases()
+        return format_tool_result(aliases) if aliases else "No aliases saved."
+    except Exception as e:
+        return log_and_format_error("list_contact_aliases", e)
+
+
+@mcp.tool(annotations=ToolAnnotations(title="Delete Contact Alias", openWorldHint=True))
+@with_account(readonly=True)
+async def delete_contact_alias(alias: str, account: Optional[str] = None) -> str:
+    """Delete a saved favorite alias."""
+    try:
+        aliases = load_aliases()
+        key = alias.strip().lstrip("@").lower()
+        if key not in aliases:
+            return f"Alias '{alias}' not found."
+        del aliases[key]
+        save_aliases(aliases)
+        return f"Alias '{alias}' deleted."
+    except Exception as e:
+        return log_and_format_error("delete_contact_alias", e, alias=alias)
+
+
 __all__ = [
     "list_contacts",
     "search_contacts",
@@ -611,4 +669,7 @@ async def send_contact(
     "export_contacts",
     "get_blocked_users",
     "send_contact",
+    "set_contact_alias",
+    "list_contact_aliases",
+    "delete_contact_alias",
 ]

tests/test_aliases.py

@@ -0,0 +1,40 @@
+"""Tests for favorite contact aliases."""
+
+from telegram_mcp import runtime
+
+
+def _use_tmp_aliases(monkeypatch, tmp_path):
+    monkeypatch.setattr(runtime, "_ALIASES_FILE", tmp_path / "aliases.json")
+
+
+def test_apply_alias_returns_saved_id(monkeypatch, tmp_path):
+    _use_tmp_aliases(monkeypatch, tmp_path)
+    runtime.save_aliases({"андрей": 12345})
+
+    assert runtime.apply_alias("андрей") == 12345
+    assert runtime.apply_alias("Андрей") == 12345
+    assert runtime.apply_alias("@андрей") == 12345
+    assert runtime.apply_alias(" андрей ") == 12345
+
+
+def test_apply_alias_passes_through_unknown_values(monkeypatch, tmp_path):
+    _use_tmp_aliases(monkeypatch, tmp_path)
+    runtime.save_aliases({"андрей": 12345})
+
+    assert runtime.apply_alias("bob") == "bob"
+    assert runtime.apply_alias(678) == 678
+
+
+def test_load_aliases_missing_or_corrupt_file(monkeypatch, tmp_path):
+    _use_tmp_aliases(monkeypatch, tmp_path)
+    assert runtime.load_aliases() == {}
+
+    (tmp_path / "aliases.json").write_text("not json")
+    assert runtime.load_aliases() == {}
+
+
+def test_save_and_load_roundtrip_lowercases_keys(monkeypatch, tmp_path):
+    _use_tmp_aliases(monkeypatch, tmp_path)
+    runtime.save_aliases({"Работа": -1001234567890})
+
+    assert runtime.load_aliases() == {"работа": -1001234567890}

tests/test_session_string_generator.py

@@ -112,6 +112,25 @@ def test_qr_login_uses_hidden_password_prompt_for_2fa(monkeypatch):
     assert client.sign_in_calls == ["secret"]
 
 
+def test_qr_login_retries_2fa_on_invalid_password(monkeypatch):
+    qr = _FakeQR([session_string_generator.errors.SessionPasswordNeededError(request=None)])
+    client = _FakeClient(qr)
+    attempts = []
+
+    def sign_in(password=None):
+        attempts.append(password)
+        if len(attempts) == 1:
+            raise session_string_generator.errors.PasswordHashInvalidError(request=None)
+
+    client.sign_in = sign_in
+    entered = iter(["wrong", "right"])
+    monkeypatch.setattr("getpass.getpass", lambda prompt: next(entered))
+
+    session_string_generator._qr_login(client)
+
+    assert attempts == ["wrong", "right"]
+
+
 def test_phone_login_uses_hidden_password_prompt_for_2fa(monkeypatch):
     class _Client:
         def __init__(self):