Merge pull request chigwell#144 from shayanline/feat/configurable-device-identity-and-qr-refresh

feat: configurable device identity and self-refreshing QR login

Author: chigwell

.env.example

@@ -14,6 +14,15 @@ TELEGRAM_SESSION_NAME=telegram_session
 # TELEGRAM_SESSION_STRING_WORK=<session string for work account>
 # TELEGRAM_SESSION_STRING_PERSONAL=<session string for personal account>
 
+# --- Device identity (optional) ---
+# Controls how this client appears in Telegram > Settings > Devices. If unset,
+# Telethon falls back to the host platform (e.g. "arm64"). Set these so the
+# session keeps a stable, recognisable name across reconnects. Applied both
+# when generating a session string and when the server connects.
+# TELEGRAM_DEVICE_MODEL=Telegram MCP
+# TELEGRAM_SYSTEM_VERSION=1.0
+# TELEGRAM_APP_VERSION=1.0
+
 # --- MCP tool exposure (optional) ---
 # Default is all. Set to read-only to expose only tools annotated with
 # readOnlyHint=True through MCP. This does not reduce Telegram session authority

README.md

@@ -30,6 +30,7 @@ Message sent successfully:
 - [Quick Start](#quick-start)
 - [MCP Client Configuration](#mcp-client-configuration)
 - [Multi-Account Setup](#multi-account-setup)
+- [Device Identity](#device-identity)
 - [Proxy Support](#proxy-support)
 - [File Path Security](#file-path-security)
 - [Docker](#docker)
@@ -222,6 +223,24 @@ Example prompts:
 - "Show unread messages from all accounts"
 - "Send this from my work account to @example"
 
+## Device Identity
+
+These optional variables control how the client appears in Telegram under
+**Settings > Devices** (the active-sessions list):
+
+```env
+TELEGRAM_DEVICE_MODEL=Telegram MCP
+TELEGRAM_SYSTEM_VERSION=1.0
+TELEGRAM_APP_VERSION=1.0
+```
+
+If left unset, Telethon falls back to the host platform (for example `arm64`).
+Because these values are re-sent on every connection, a long-running server
+would otherwise overwrite the name chosen during login on each reconnect, so
+set them to keep a stable, recognisable device name. The same variables are
+read both by the session string generator (at login) and by the server (on
+every connect), so set them in the same place as your other credentials.
+
 ## Proxy Support
 
 Route Telegram traffic through a proxy by setting the `TELEGRAM_PROXY_*`

session_string_generator.py

@@ -26,13 +26,18 @@
 import io
 import os
 import sys
+from datetime import datetime, timezone
 
 from dotenv import load_dotenv
 from telethon import errors
 from telethon.sessions import StringSession
 from telethon.sync import TelegramClient
+from telegram_mcp.client_identity import client_identity_kwargs
 from telegram_mcp.install_guard import UnsafeInstallationError, assert_safe_distribution
 
+# How many times the QR code is regenerated after expiry before giving up.
+_QR_MAX_REFRESHES = 10
+
 load_dotenv()
 
 
@@ -62,11 +67,9 @@ def _check_installation() -> None:
         sys.exit(1)
 
 
-def _qr_login(client: TelegramClient) -> None:
+def _render_qr(qr) -> None:
     import qrcode
 
-    qr = client.qr_login()
-
     print("\n----- QR Code Login -----\n")
 
     qr_obj = qrcode.QRCode(border=1)
@@ -82,15 +85,38 @@ def _qr_login(client: TelegramClient) -> None:
     print(f"Expires at: {qr.expires.strftime('%H:%M:%S')}")
     print("Waiting for you to scan...")
 
-    try:
-        client.loop.run_until_complete(qr.wait(timeout=120))
-    except asyncio.TimeoutError:
-        print("\nQR code expired. Please try again.")
-        client.disconnect()
-        sys.exit(1)
-    except errors.SessionPasswordNeededError:
-        pw = getpass.getpass("\nTwo-factor authentication enabled. Please enter your password: ")
-        client.sign_in(password=pw)
+
+def _seconds_until_expiry(qr) -> float:
+    """Seconds left before this QR token expires, with a small safety margin."""
+    expires = qr.expires
+    if expires.tzinfo is None:
+        expires = expires.replace(tzinfo=timezone.utc)
+    remaining = (expires - datetime.now(timezone.utc)).total_seconds()
+    return max(1.0, remaining - 1.0)
+
+
+def _qr_login(client: TelegramClient) -> None:
+    qr = client.qr_login()
+    _render_qr(qr)
+
+    for _ in range(_QR_MAX_REFRESHES):
+        try:
+            client.loop.run_until_complete(qr.wait(timeout=_seconds_until_expiry(qr)))
+            return
+        except asyncio.TimeoutError:
+            client.loop.run_until_complete(qr.recreate())
+            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
+
+    print("\nQR code expired too many times. Please run the generator again.")
+    client.disconnect()
+    sys.exit(1)
 
 
 def _phone_login(client: TelegramClient) -> None:
@@ -161,7 +187,7 @@ def main() -> None:
         method = input("\nEnter 1 or 2 [default: 1]: ").strip() or "1"
 
     try:
-        client = TelegramClient(StringSession(), API_ID, API_HASH)
+        client = TelegramClient(StringSession(), API_ID, API_HASH, **client_identity_kwargs())
         client.connect()
 
         if not client.is_user_authorized():

telegram_mcp/client_identity.py

@@ -0,0 +1,31 @@
+"""Resolve the Telethon device identity used when connecting to Telegram.
+
+These values are what Telegram shows in Settings > Devices (active sessions).
+When they are not provided, Telethon falls back to the host platform (for
+example ``arm64``), and because the values are re-sent on every connection,
+a long-running server would otherwise overwrite the name chosen during login
+on each reconnect. Making them configurable keeps a stable, recognisable name.
+"""
+
+import os
+
+# Maps the Telethon constructor keyword to the environment variable that sets it.
+_DEVICE_ENV = {
+    "device_model": "TELEGRAM_DEVICE_MODEL",
+    "system_version": "TELEGRAM_SYSTEM_VERSION",
+    "app_version": "TELEGRAM_APP_VERSION",
+}
+
+
+def client_identity_kwargs() -> dict:
+    """Return TelegramClient device kwargs derived from the environment.
+
+    Only variables that are set to a non-empty value are included, so unset
+    ones keep Telethon's own defaults.
+    """
+    kwargs = {}
+    for kwarg, env_var in _DEVICE_ENV.items():
+        value = os.environ.get(env_var)
+        if value:
+            kwargs[kwarg] = value
+    return kwargs

telegram_mcp/runtime.py

@@ -45,6 +45,7 @@
 from functools import wraps
 import telethon.errors.rpcerrorlist
 from sanitize import sanitize_user_content, sanitize_name, sanitize_dict, format_tool_result
+from telegram_mcp.client_identity import client_identity_kwargs
 
 
 class ValidationError(Exception):
@@ -275,6 +276,7 @@ def _build_client(session: Any, label: str) -> TelegramClient:
         kwargs["proxy"] = proxy
     if connection is not None:
         kwargs["connection"] = connection
+    kwargs.update(client_identity_kwargs())
     return TelegramClient(session, TELEGRAM_API_ID, TELEGRAM_API_HASH, **kwargs)
 
 

tests/test_client_identity.py

@@ -0,0 +1,33 @@
+from telegram_mcp.client_identity import client_identity_kwargs
+
+
+def _clear_device_env(monkeypatch):
+    for key in ("TELEGRAM_DEVICE_MODEL", "TELEGRAM_SYSTEM_VERSION", "TELEGRAM_APP_VERSION"):
+        monkeypatch.delenv(key, raising=False)
+
+
+def test_client_identity_kwargs_empty_when_unset(monkeypatch):
+    _clear_device_env(monkeypatch)
+
+    assert client_identity_kwargs() == {}
+
+
+def test_client_identity_kwargs_maps_all_variables(monkeypatch):
+    _clear_device_env(monkeypatch)
+    monkeypatch.setenv("TELEGRAM_DEVICE_MODEL", "Telegram MCP")
+    monkeypatch.setenv("TELEGRAM_SYSTEM_VERSION", "2.0")
+    monkeypatch.setenv("TELEGRAM_APP_VERSION", "3.1")
+
+    assert client_identity_kwargs() == {
+        "device_model": "Telegram MCP",
+        "system_version": "2.0",
+        "app_version": "3.1",
+    }
+
+
+def test_client_identity_kwargs_ignores_empty_values(monkeypatch):
+    _clear_device_env(monkeypatch)
+    monkeypatch.setenv("TELEGRAM_DEVICE_MODEL", "Telegram MCP")
+    monkeypatch.setenv("TELEGRAM_SYSTEM_VERSION", "")
+
+    assert client_identity_kwargs() == {"device_model": "Telegram MCP"}

tests/test_runtime.py

@@ -257,6 +257,25 @@ def test_discover_accounts_passes_proxy_kwargs_to_client(monkeypatch):
     assert client.kwargs["connection"] is ConnectionTcpMTProxyRandomizedIntermediate
 
 
+def test_discover_accounts_passes_device_identity_kwargs_to_client(monkeypatch):
+    _clear_session_env(monkeypatch)
+    _clear_proxy_env(monkeypatch)
+    for key in ("TELEGRAM_DEVICE_MODEL", "TELEGRAM_SYSTEM_VERSION", "TELEGRAM_APP_VERSION"):
+        monkeypatch.delenv(key, raising=False)
+    monkeypatch.setenv("TELEGRAM_SESSION_STRING", "default-session")
+    monkeypatch.setenv("TELEGRAM_DEVICE_MODEL", "Telegram MCP")
+    monkeypatch.setenv("TELEGRAM_APP_VERSION", "3.1")
+    monkeypatch.setattr(runtime, "TelegramClient", _FakeTelegramClient)
+    monkeypatch.setattr(runtime, "StringSession", lambda value: f"StringSession:{value}")
+
+    accounts = runtime._discover_accounts()
+
+    client = accounts["default"]
+    assert client.kwargs["device_model"] == "Telegram MCP"
+    assert client.kwargs["app_version"] == "3.1"
+    assert "system_version" not in client.kwargs
+
+
 def test_get_client_single_and_multi_account_paths(monkeypatch):
     only = object()
     monkeypatch.setattr(runtime, "clients", {"only": only})

tests/test_session_string_generator.py

@@ -1,55 +1,110 @@
+import asyncio
+from datetime import datetime, timedelta, timezone
+
 import pytest
 
 import session_string_generator
 
 
-def test_qr_login_uses_hidden_password_prompt_for_2fa(monkeypatch):
-    class _Loop:
-        def run_until_complete(self, awaitable):
-            awaitable.close()
-            raise session_string_generator.errors.SessionPasswordNeededError(request=None)
+class _NS:
+    def __init__(self, expires):
+        self.expires = expires
 
-    class _Client:
-        def __init__(self):
-            self.loop = _Loop()
-            self.sign_in_calls = []
 
-        def qr_login(self):
-            class _Qr:
-                url = "https://example.test/qr"
+class _FakeQR:
+    def __init__(self, wait_outcomes):
+        # Each entry is either an exception to raise on wait(), or None to succeed.
+        self._wait_outcomes = list(wait_outcomes)
+        self.url = "tg://login?token=fake"
+        self.expires = datetime.now(timezone.utc) + timedelta(seconds=30)
+        self.recreate_calls = 0
 
-                class _Expiry:
-                    def strftime(self, _fmt):
-                        return "00:00:00"
+    def wait(self, timeout=None):
+        return ("wait", self._wait_outcomes.pop(0))
 
-                expires = _Expiry()
+    def recreate(self):
+        self.recreate_calls += 1
+        self.expires = datetime.now(timezone.utc) + timedelta(seconds=30)
+        return ("recreate", None)
 
-                async def wait(self, timeout=120):
-                    return None
 
-            return _Qr()
+class _FakeLoop:
+    def run_until_complete(self, marker):
+        kind, outcome = marker
+        if kind == "wait" and outcome is not None:
+            raise outcome
+        return None
 
-        def sign_in(self, password=None):
-            self.sign_in_calls.append(password)
 
-    class _FakeQrCode:
-        def __init__(self, border=1):
-            self.border = border
+class _FakeClient:
+    def __init__(self, qr):
+        self._qr = qr
+        self.loop = _FakeLoop()
+        self.disconnected = False
+        self.sign_in_calls = []
 
-        def add_data(self, _data):
-            return None
+    def qr_login(self):
+        return self._qr
 
-        def make(self, fit=True):
-            return None
+    def sign_in(self, password=None):
+        self.sign_in_calls.append(password)
 
-        def print_ascii(self, out, invert=True):
-            out.write("qr")
+    def disconnect(self):
+        self.disconnected = True
 
-    client = _Client()
+
+def test_seconds_until_expiry_handles_naive_and_aware():
+    naive = datetime(2999, 1, 1)
+    aware = datetime.now(timezone.utc) + timedelta(seconds=30)
+
+    assert session_string_generator._seconds_until_expiry(_NS(naive)) > 1.0
+    assert session_string_generator._seconds_until_expiry(_NS(aware)) > 1.0
+
+
+def test_seconds_until_expiry_never_below_one():
+    past = datetime.now(timezone.utc) - timedelta(seconds=10)
+
+    assert session_string_generator._seconds_until_expiry(_NS(past)) == 1.0
+
+
+def test_qr_login_returns_on_first_successful_scan():
+    qr = _FakeQR([None])
+    client = _FakeClient(qr)
+
+    session_string_generator._qr_login(client)
+
+    assert qr.recreate_calls == 0
+    assert client.disconnected is False
+
+
+def test_qr_login_refreshes_after_expiry_then_succeeds():
+    qr = _FakeQR([asyncio.TimeoutError(), None])
+    client = _FakeClient(qr)
+
+    session_string_generator._qr_login(client)
+
+    assert qr.recreate_calls == 1
+    assert client.disconnected is False
+
+
+def test_qr_login_gives_up_after_max_refreshes(monkeypatch):
+    monkeypatch.setattr(session_string_generator, "_QR_MAX_REFRESHES", 3)
+    qr = _FakeQR([asyncio.TimeoutError()] * 3)
+    client = _FakeClient(qr)
+
+    with pytest.raises(SystemExit):
+        session_string_generator._qr_login(client)
+
+    assert qr.recreate_calls == 3
+    assert client.disconnected is True
+
+
+def test_qr_login_uses_hidden_password_prompt_for_2fa(monkeypatch):
+    qr = _FakeQR([session_string_generator.errors.SessionPasswordNeededError(request=None)])
+    client = _FakeClient(qr)
     prompted = []
 
     monkeypatch.setattr("getpass.getpass", lambda prompt: prompted.append(prompt) or "secret")
-    monkeypatch.setattr("qrcode.QRCode", _FakeQrCode)
 
     session_string_generator._qr_login(client)