Merge pull request chigwell#259 from vsem-azamat/fix/non-interactive-reconnect

fix(runtime): don't start interactive login on reconnect; add optional TELEGRAM_EXPECTED_USERNAME check

Author: chigwell

.env.example

@@ -50,6 +50,11 @@ TELEGRAM_SESSION_NAME=telegram_session
 # TELEGRAM_SYSTEM_VERSION=1.0
 # TELEGRAM_APP_VERSION=1.0
 
+# --- Expected account check (optional) ---
+# Refuse to start if the session belongs to another account. _<LABEL> overrides.
+# TELEGRAM_EXPECTED_USERNAME=my_username
+# TELEGRAM_EXPECTED_USERNAME_WORK=my_work_username
+
 # --- 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

@@ -32,6 +32,7 @@ Message sent successfully:
 - [MCP Client Configuration](#mcp-client-configuration)
 - [Multi-Account Setup](#multi-account-setup)
 - [Device Identity](#device-identity)
+- [Expected Account Check](#expected-account-check)
 - [Proxy Support](#proxy-support)
 - [File Path Security](#file-path-security)
 - [Chat Access Privacy (Allowlist)](#chat-access-privacy-allowlist)
@@ -567,6 +568,10 @@ 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.
 
+## Expected Account Check
+
+Set `TELEGRAM_EXPECTED_USERNAME` (or `TELEGRAM_EXPECTED_USERNAME_<LABEL>` per account) and the server refuses to start when the session belongs to a different account. The comparison is case-insensitive and ignores a leading `@`. Unset means no check.
+
 ## Proxy Support
 
 Route Telegram traffic through a proxy by setting the `TELEGRAM_PROXY_*`
@@ -842,7 +847,8 @@ Telegram messages, display names, chat titles, and button labels are untrusted c
   the MCP server when you can scan from an existing Telegram app, or
   `uv run session_string_generator.py --phone` when you need phone-code login.
   Then set `TELEGRAM_SESSION_STRING` in `.env`. The MCP server does not perform
-  interactive phone-code login over stdio.
+  interactive phone-code login over stdio, including after a reconnect: a session
+  terminated while the server runs fails with "not authorized after reconnect".
 - **Invalid API credentials:** verify `TELEGRAM_API_ID` and `TELEGRAM_API_HASH` at [my.telegram.org/apps](https://my.telegram.org/apps).
 - **Database is locked:** make sure no other process is using the same file session.
 - **`AuthKeyDuplicatedError` / "Another telegram-mcp process is already connected with this session":** two processes tried to connect the same Telegram session at once (e.g. an MCP client restarted the connector before the old process exited), which Telegram rejects and can invalidate the session for both. The server now takes an exclusive lock per session before connecting; a second concurrent launch waits briefly (default 20s, override with `TELEGRAM_LOCK_GRACE_SECONDS`) for the first to release it and otherwise exits without ever calling `connect()`, instead of racing into a duplicate connection. Retry once only one instance is running — the refusal names the PID holding the lock. If several instances on this host are meant to share one session (all reaching Telegram from the same IP), set `TELEGRAM_SESSION_LOCK=shared`; see [Sharing one session from one host](#sharing-one-session-from-one-host).

telegram_mcp/runner.py

@@ -63,6 +63,44 @@ def _session_lock_shared() -> bool:
     return raw == "shared"
 
 
+def _normalize_username(value) -> str:
+    return (value or "").strip().lstrip("@").casefold()
+
+
+def _expected_username(label: str) -> str:
+    """TELEGRAM_EXPECTED_USERNAME_<LABEL>, else TELEGRAM_EXPECTED_USERNAME; "" if unset."""
+    raw = os.getenv(f"TELEGRAM_EXPECTED_USERNAME_{label.upper()}") or os.getenv(
+        "TELEGRAM_EXPECTED_USERNAME"
+    )
+    return _normalize_username(raw)
+
+
+async def _verify_expected_username(label: str, client) -> None:
+    """Refuse to serve a session logged in to a different account."""
+    expected = _expected_username(label)
+    if not expected:
+        return
+
+    me = await client.get_me()
+    usernames = {_normalize_username(getattr(me, "username", None))}
+    for entry in getattr(me, "usernames", None) or []:
+        if getattr(entry, "active", False):
+            usernames.add(_normalize_username(getattr(entry, "username", None)))
+    usernames.discard("")
+
+    if expected in usernames:
+        return
+
+    try:
+        await client.disconnect()
+    except Exception:
+        pass
+    raise RuntimeError(
+        f"Telegram client '{label}' is logged in to a different account than "
+        f"TELEGRAM_EXPECTED_USERNAME_{label.upper()} / TELEGRAM_EXPECTED_USERNAME expects."
+    )
+
+
 async def _connect_authorized_client(label, client) -> None:
     # First, prevent our own duplicate-spawn case outright: a per-session lock
     # means a second instance of this server never even attempts to connect
@@ -106,6 +144,7 @@ async def _connect_authorized_client(label, client) -> None:
             await asyncio.sleep(delay)
 
     if await client.is_user_authorized():
+        await _verify_expected_username(label, client)
         return
 
     raise RuntimeError(

telegram_mcp/runtime.py

@@ -813,8 +813,16 @@ async def _force_reconnect(cl: TelegramClient):
             f"Reconnecting to Telegram timed out after {_RECONNECT_TIMEOUT:.0f}s."
         ) from exc
     if not await cl.is_user_authorized():
-        reconnect_logger.warning("Client not authorized after reconnect, calling start()...")
-        await asyncio.wait_for(cl.start(), timeout=_RECONNECT_TIMEOUT)
+        # cl.start() would prompt via blocking input(): it stalls the event loop
+        # and, over stdio, reads protocol frames as a phone number.
+        try:
+            await cl.disconnect()
+        except Exception:
+            pass
+        raise RuntimeError(
+            "Telegram session is not authorized after reconnect. Re-authorize it "
+            "outside the server (session_string_generator.py) and restart."
+        )
     _last_conn_verified[id(cl)] = time.time()
     reconnect_logger.warning("Forced reconnect successful")
 

tests/conftest.py

@@ -15,3 +15,11 @@ def transcript_cache_dir(tmp_path, monkeypatch):
     d = tmp_path / "transcripts"
     monkeypatch.setenv("TELEGRAM_TRANSCRIPT_CACHE_DIR", str(d))
     return d
+
+
+@pytest.fixture(autouse=True)
+def _clear_expected_username(monkeypatch):
+    """Keep a developer's TELEGRAM_EXPECTED_USERNAME* (.env) away from fake clients."""
+    for key in list(os.environ):
+        if key.startswith("TELEGRAM_EXPECTED_USERNAME"):
+            monkeypatch.delenv(key, raising=False)

tests/test_runner.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from types import SimpleNamespace
 
 import pytest
 
@@ -16,18 +17,36 @@ def save(self):
 
 
 class _FakeClient:
-    def __init__(self, *, authorized: bool, identity: str = "test-identity"):
+    def __init__(
+        self,
+        *,
+        authorized: bool,
+        identity: str = "test-identity",
+        username: str | None = "work_account",
+        usernames: list | None = None,
+    ):
         self.authorized = authorized
         self.connected = False
+        self.disconnected = False
         self.started = False
+        self.get_me_calls = 0
+        self.username = username
+        self.usernames = usernames
         self.session = _FakeSession(identity)
 
     async def connect(self):
         self.connected = True
 
+    async def disconnect(self):
+        self.disconnected = True
+
     async def is_user_authorized(self):
         return self.authorized
 
+    async def get_me(self):
+        self.get_me_calls += 1
+        return SimpleNamespace(username=self.username, usernames=self.usernames)
+
     async def start(self):
         self.started = True
 
@@ -75,6 +94,110 @@ async def test_connect_authorized_client_rejects_unauthorized_session():
     assert client.started is False
 
 
+@pytest.mark.asyncio
+async def test_connect_authorized_client_skips_identity_check_when_unset():
+    client = _FakeClient(authorized=True, username="anyone")
+
+    await runner._connect_authorized_client("default", client)
+
+    assert client.get_me_calls == 0
+    assert client.disconnected is False
+
+
+@pytest.mark.asyncio
+async def test_connect_authorized_client_accepts_expected_username(monkeypatch):
+    monkeypatch.setenv("TELEGRAM_EXPECTED_USERNAME", " @Work_Account ")
+    client = _FakeClient(authorized=True, username="work_account")
+
+    await runner._connect_authorized_client("default", client)
+
+    assert client.get_me_calls == 1
+    assert client.disconnected is False
+
+
+@pytest.mark.asyncio
+async def test_connect_authorized_client_rejects_wrong_account(monkeypatch):
+    monkeypatch.setenv("TELEGRAM_EXPECTED_USERNAME", "work_account")
+    client = _FakeClient(authorized=True, username="personal_account")
+
+    with pytest.raises(RuntimeError, match="different account") as excinfo:
+        await runner._connect_authorized_client("default", client)
+
+    assert "TELEGRAM_EXPECTED_USERNAME" in str(excinfo.value)
+    # The other account's username is not echoed into logs.
+    assert "personal_account" not in str(excinfo.value)
+    assert client.disconnected is True
+
+
+@pytest.mark.asyncio
+async def test_connect_authorized_client_rejects_account_without_username(monkeypatch):
+    monkeypatch.setenv("TELEGRAM_EXPECTED_USERNAME", "work_account")
+    client = _FakeClient(authorized=True, username=None)
+
+    with pytest.raises(RuntimeError, match="different account"):
+        await runner._connect_authorized_client("default", client)
+
+    assert client.disconnected is True
+
+
+@pytest.mark.asyncio
+async def test_connect_authorized_client_accepts_active_collectible_username(monkeypatch):
+    monkeypatch.setenv("TELEGRAM_EXPECTED_USERNAME", "collectible")
+    client = _FakeClient(
+        authorized=True,
+        username=None,
+        usernames=[
+            SimpleNamespace(username="inactive", active=False),
+            SimpleNamespace(username="Collectible", active=True),
+        ],
+    )
+
+    await runner._connect_authorized_client("default", client)
+
+    assert client.disconnected is False
+
+
+@pytest.mark.asyncio
+async def test_connect_authorized_client_ignores_inactive_collectible_username(monkeypatch):
+    monkeypatch.setenv("TELEGRAM_EXPECTED_USERNAME", "inactive")
+    client = _FakeClient(
+        authorized=True,
+        username="primary",
+        usernames=[SimpleNamespace(username="inactive", active=False)],
+    )
+
+    with pytest.raises(RuntimeError, match="different account"):
+        await runner._connect_authorized_client("default", client)
+
+
+@pytest.mark.asyncio
+async def test_connect_authorized_client_per_label_expected_username(monkeypatch):
+    monkeypatch.setenv("TELEGRAM_EXPECTED_USERNAME", "personal_account")
+    monkeypatch.setenv("TELEGRAM_EXPECTED_USERNAME_WORK", "work_account")
+    work = _FakeClient(authorized=True, identity="session-work", username="work_account")
+    personal = _FakeClient(
+        authorized=True, identity="session-personal", username="personal_account"
+    )
+
+    await runner._connect_authorized_client("work", work)
+    await runner._connect_authorized_client("personal", personal)
+
+    assert work.disconnected is False
+    assert personal.disconnected is False
+
+
+@pytest.mark.asyncio
+async def test_connect_authorized_client_per_label_mismatch_names_label(monkeypatch):
+    monkeypatch.setenv("TELEGRAM_EXPECTED_USERNAME_WORK", "work_account")
+    client = _FakeClient(authorized=True, username="personal_account")
+
+    with pytest.raises(RuntimeError, match="TELEGRAM_EXPECTED_USERNAME_WORK") as excinfo:
+        await runner._connect_authorized_client("work", client)
+
+    assert "'work'" in str(excinfo.value)
+    assert client.disconnected is True
+
+
 @pytest.mark.asyncio
 async def test_connect_authorized_client_refuses_concurrent_duplicate_session():
     first = _FakeClient(authorized=True, identity="shared-session")

tests/test_runtime.py

@@ -545,15 +545,42 @@ async def __call__(self, _request):
 
 @pytest.mark.asyncio
 async def test_ensure_connected_reconnects_disconnected_client(monkeypatch):
-    client = _ConnectivityClient(connected=False, authorized=False)
+    client = _ConnectivityClient(connected=False, authorized=True)
     monkeypatch.setattr(runtime, "_last_conn_verified", {})
 
     await runtime.ensure_connected(client)
 
-    assert client.calls == ["is_connected", "disconnect", "connect", "is_user_authorized", "start"]
+    assert client.calls == ["is_connected", "disconnect", "connect", "is_user_authorized"]
     assert runtime._last_conn_verified[id(client)] > 0
 
 
+@pytest.mark.asyncio
+async def test_force_reconnect_never_starts_interactive_login(monkeypatch):
+    client = _ConnectivityClient(connected=False, authorized=False)
+    monkeypatch.setattr(runtime, "_last_conn_verified", {})
+
+    with pytest.raises(RuntimeError, match="not authorized after reconnect") as excinfo:
+        await runtime._force_reconnect(client)
+
+    assert "session_string_generator.py" in str(excinfo.value)
+    assert "start" not in client.calls
+    # The unauthorized connection is dropped so the next tool call reconnects
+    # and reports the same error instead of reusing a half-open session.
+    assert client.calls == ["disconnect", "connect", "is_user_authorized", "disconnect"]
+    assert id(client) not in runtime._last_conn_verified
+
+
+@pytest.mark.asyncio
+async def test_ensure_connected_surfaces_revoked_session(monkeypatch):
+    client = _ConnectivityClient(connected=False, authorized=False)
+    monkeypatch.setattr(runtime, "_last_conn_verified", {})
+
+    with pytest.raises(RuntimeError, match="not authorized after reconnect"):
+        await runtime.ensure_connected(client)
+
+    assert "start" not in client.calls
+
+
 @pytest.mark.asyncio
 async def test_ensure_connected_pings_and_reconnects_on_failed_ping(monkeypatch):
     client = _ConnectivityClient(connected=True, authorized=True, ping_error=ConnectionError())