Merge pull request #152 from nikit34/feat/session-pool-authkey

feat: session pool + AuthKeyDuplicatedError retry for concurrent clients on one account

Author: chigwell

.env.example

@@ -14,6 +14,17 @@ TELEGRAM_SESSION_NAME=telegram_session
 # TELEGRAM_SESSION_STRING_WORK=<session string for work account>
 # TELEGRAM_SESSION_STRING_PERSONAL=<session string for personal account>
 
+# --- Session pool (same account, concurrent clients) ---
+# Run several MCP clients (e.g. the desktop app AND a terminal CLI) against ONE
+# Telegram account without hitting AuthKeyDuplicatedError. Telegram forbids a
+# single session (auth key) being used from two IPs at once; on a VPN/dual-stack
+# host two local clients can collide. List several interchangeable session
+# strings (whitespace, comma or semicolon separated) and each process claims a
+# free one via an advisory file lock. Generate extra sessions with
+# `uv run session_string_generator.py`. Takes precedence over
+# TELEGRAM_SESSION_STRING for the default account.
+# TELEGRAM_SESSION_STRINGS=<session A> <session B> <session C>
+
 # --- 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

README.md

@@ -263,6 +263,26 @@ Example prompts:
 
 - "List my accounts"
 - "Show unread messages from all accounts"
+
+### Session pool (one account, several concurrent clients)
+
+To run several MCP clients against the **same** Telegram account at once (for
+example the desktop app *and* a terminal CLI), give each client its own
+authorized session. Telegram forbids one session (auth key) being used from two
+IPs simultaneously, so on a VPN or dual-stack host two local clients can collide
+with `AuthKeyDuplicatedError`. List several interchangeable session strings in
+`TELEGRAM_SESSION_STRINGS` (separated by whitespace, comma or semicolon); each
+process claims a free one via an advisory file lock, so clients deterministically
+pick distinct sessions:
+
+```env
+TELEGRAM_SESSION_STRINGS=<session A> <session B> <session C>
+```
+
+Generate extra sessions with `uv run session_string_generator.py`. The pool
+takes precedence over `TELEGRAM_SESSION_STRING` for the default account. As an
+extra safety net, a transient `AuthKeyDuplicatedError` at connect time (e.g.
+during a VPN reconnect) is retried with backoff before the server gives up.
 - "Send this from my work account to @example"
 
 ## Device Identity

telegram_mcp/runner.py

@@ -7,13 +7,42 @@
 except UnsafeInstallationError as exc:
     raise SystemExit(str(exc)) from None
 
+from telethon.errors import AuthKeyDuplicatedError
+
 from telegram_mcp import runtime as _runtime
 from telegram_mcp.runtime import *
 import telegram_mcp.tools  # noqa: F401 - registers MCP tools via decorators
 
 
 async def _connect_authorized_client(label, client) -> None:
-    await client.connect()
+    # Tolerate a transient AuthKeyDuplicatedError (the same session briefly seen
+    # from two IPs, e.g. during a VPN reconnect) with a bounded retry so a blip
+    # does not take the whole server down. Give each concurrent client its own
+    # session (TELEGRAM_SESSION_STRINGS pool or TELEGRAM_SESSION_STRING_<LABEL>)
+    # to avoid the collision entirely.
+    max_attempts = 4
+    for attempt in range(1, max_attempts + 1):
+        try:
+            await client.connect()
+            break
+        except AuthKeyDuplicatedError:
+            if attempt >= max_attempts:
+                raise
+            delay = min(2**attempt, 15)
+            print(
+                f"AuthKeyDuplicatedError connecting '{label}' (attempt "
+                f"{attempt}/{max_attempts}): session in use from another IP. "
+                f"Retrying in {delay}s. If this persists, give each concurrent "
+                "client its own session via TELEGRAM_SESSION_STRINGS or "
+                "TELEGRAM_SESSION_STRING_<LABEL>.",
+                file=sys.stderr,
+            )
+            try:
+                await client.disconnect()
+            except Exception:
+                pass
+            await asyncio.sleep(delay)
+
     if await client.is_user_authorized():
         return
 

telegram_mcp/runtime.py

@@ -42,6 +42,14 @@
     TextWithEntities,
 )
 import re
+import hashlib
+import tempfile
+
+try:
+    import fcntl  # POSIX advisory locks; unavailable on Windows
+except ImportError:  # pragma: no cover - Windows fallback
+    fcntl = None
+
 from functools import wraps
 import telethon.errors.rpcerrorlist
 from sanitize import sanitize_user_content, sanitize_name, sanitize_dict, format_tool_result
@@ -280,11 +288,88 @@ def _build_client(session: Any, label: str) -> TelegramClient:
     return TelegramClient(session, TELEGRAM_API_ID, TELEGRAM_API_HASH, **kwargs)
 
 
+# --- Session pool ------------------------------------------------------------
+# A POOL of interchangeable authorized sessions for the SAME account lets
+# several concurrent MCP clients (e.g. the desktop app AND a terminal CLI) run
+# against one Telegram account without tripping AuthKeyDuplicatedError.
+#
+# Telegram forbids one auth key (one StringSession) being used from two IPs at
+# once; on a dual-stack / VPN host two local clients can egress via different
+# source IPs and collide. The fix is one authorized session PER concurrent
+# client (Telegram allows one account on many "devices"). Generate extra
+# sessions with `uv run session_string_generator.py` and list them in
+# TELEGRAM_SESSION_STRINGS (whitespace/comma/semicolon separated). Each process
+# claims the first session not already locked by a live process via an advisory
+# flock, so clients deterministically pick distinct slots; the OS releases the
+# lock if a process dies.
+
+# Acquired lock handles are held for the process lifetime so the advisory locks
+# stay held until exit (or crash, when the OS releases them).
+_SESSION_LOCKS: list = []
+
+
+def _parse_session_pool() -> List[str]:
+    """Parse TELEGRAM_SESSION_STRINGS into a de-duplicated list of sessions."""
+    raw = os.getenv("TELEGRAM_SESSION_STRINGS")
+    if not raw:
+        return []
+    pool: List[str] = []
+    for tok in re.split(r"[\s,;]+", raw.strip()):
+        if tok and tok not in pool:
+            pool.append(tok)
+    return pool
+
+
+def _acquire_session(pool: List[str]) -> str:
+    """Claim the first free session in the pool via an advisory file lock."""
+    if fcntl is None:
+        # No advisory locks (e.g. Windows): can't coordinate slots, so use the
+        # first session. For concurrent clients there, prefer distinct
+        # TELEGRAM_SESSION_STRING_<LABEL> accounts instead.
+        return pool[0]
+    lock_dir = os.path.join(tempfile.gettempdir(), "telegram-mcp-session-locks")
+    try:
+        os.makedirs(lock_dir, exist_ok=True)
+    except OSError:
+        lock_dir = tempfile.gettempdir()
+    for idx, session in enumerate(pool):
+        digest = hashlib.sha1(session.encode("utf-8")).hexdigest()[:16]
+        lock_path = os.path.join(lock_dir, f"session-{digest}.lock")
+        try:
+            fh = open(lock_path, "w")
+            fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except OSError:
+            # Locked by another live client — try the next session.
+            try:
+                fh.close()
+            except Exception:
+                pass
+            continue
+        _SESSION_LOCKS.append(fh)
+        try:
+            fh.write(f"pid={os.getpid()}\n")
+            fh.flush()
+        except OSError:
+            pass
+        print(f"Using Telegram session slot {idx + 1}/{len(pool)}.", file=sys.stderr)
+        return session
+    print(
+        f"WARNING: all {len(pool)} pooled Telegram session(s) are already in use "
+        "by other clients; reusing the first (may raise AuthKeyDuplicatedError). "
+        "Add another session to TELEGRAM_SESSION_STRINGS to run more clients.",
+        file=sys.stderr,
+    )
+    return pool[0]
+
+
 def _discover_accounts() -> dict[str, TelegramClient]:
     """Scan env vars to build account label -> TelegramClient mapping.
 
     Detection rules:
     - TELEGRAM_SESSION_STRING_<LABEL> / TELEGRAM_SESSION_NAME_<LABEL> -> multi-mode
+    - TELEGRAM_SESSION_STRINGS (whitespace/comma/semicolon separated) -> a pool
+      of interchangeable sessions for the default account; each process claims a
+      free slot to avoid AuthKeyDuplicatedError (takes precedence for "default")
     - Unsuffixed TELEGRAM_SESSION_STRING / TELEGRAM_SESSION_NAME -> label "default"
     - If both suffixed and unsuffixed exist -> unsuffixed becomes "default"
 
@@ -304,14 +389,21 @@ def _discover_accounts() -> dict[str, TelegramClient]:
             label = key[len(prefix_name) :].lower()
             accounts[label] = _build_client(value, label)
 
-    # Backward-compatible unsuffixed variables
+    # Backward-compatible unsuffixed variables. A pool (TELEGRAM_SESSION_STRINGS)
+    # takes precedence for the default account and claims a free session slot.
+    session_pool = _parse_session_pool()
     session_string = os.getenv("TELEGRAM_SESSION_STRING")
     session_name = os.getenv("TELEGRAM_SESSION_NAME")
 
-    if session_string and "default" not in accounts:
-        accounts["default"] = _build_client(StringSession(session_string), "default")
-    elif session_name and "default" not in accounts:
-        accounts["default"] = _build_client(session_name, "default")
+    if "default" not in accounts:
+        if session_pool:
+            accounts["default"] = _build_client(
+                StringSession(_acquire_session(session_pool)), "default"
+            )
+        elif session_string:
+            accounts["default"] = _build_client(StringSession(session_string), "default")
+        elif session_name:
+            accounts["default"] = _build_client(session_name, "default")
 
     if not accounts:
         print(

tests/test_session_pool.py

@@ -0,0 +1,143 @@
+import hashlib
+import os
+
+import pytest
+
+from telegram_mcp import runner, runtime
+
+# --- _parse_session_pool -----------------------------------------------------
+
+
+def test_parse_session_pool_splits_and_dedupes(monkeypatch):
+    monkeypatch.setenv("TELEGRAM_SESSION_STRINGS", "s1, s2 ;s3\n s1  s2")
+    assert runtime._parse_session_pool() == ["s1", "s2", "s3"]
+
+
+def test_parse_session_pool_empty_when_unset(monkeypatch):
+    monkeypatch.delenv("TELEGRAM_SESSION_STRINGS", raising=False)
+    assert runtime._parse_session_pool() == []
+
+
+# --- _acquire_session --------------------------------------------------------
+
+
+@pytest.fixture
+def isolated_lock_dir(tmp_path, monkeypatch):
+    monkeypatch.setattr(runtime.tempfile, "gettempdir", lambda: str(tmp_path))
+    monkeypatch.setattr(runtime, "_SESSION_LOCKS", [])
+    return tmp_path
+
+
+def _lock_slot(lock_dir, session):
+    """Simulate another live client holding the slot for ``session``."""
+    import fcntl
+
+    digest = hashlib.sha1(session.encode("utf-8")).hexdigest()[:16]
+    path = os.path.join(str(lock_dir), "telegram-mcp-session-locks", f"session-{digest}.lock")
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    fh = open(path, "w")
+    fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+    return fh
+
+
+def test_acquire_session_claims_first_free_slot(isolated_lock_dir):
+    assert runtime._acquire_session(["AAA", "BBB", "CCC"]) == "AAA"
+
+
+def test_acquire_session_skips_slot_locked_by_another_client(isolated_lock_dir):
+    foreign = _lock_slot(isolated_lock_dir, "AAA")
+    try:
+        assert runtime._acquire_session(["AAA", "BBB", "CCC"]) == "BBB"
+    finally:
+        foreign.close()
+
+
+def test_acquire_session_falls_back_to_first_when_pool_exhausted(isolated_lock_dir):
+    held = [_lock_slot(isolated_lock_dir, s) for s in ("AAA", "BBB")]
+    try:
+        # Only two slots exist and both are taken -> reuse the first.
+        assert runtime._acquire_session(["AAA", "BBB"]) == "AAA"
+    finally:
+        for fh in held:
+            fh.close()
+
+
+def test_acquire_session_without_fcntl_uses_first(isolated_lock_dir, monkeypatch):
+    monkeypatch.setattr(runtime, "fcntl", None)
+    assert runtime._acquire_session(["AAA", "BBB"]) == "AAA"
+
+
+# --- _discover_accounts prefers the pool for the default account -------------
+
+
+def test_discover_accounts_uses_pool_for_default(monkeypatch):
+    monkeypatch.setenv("TELEGRAM_SESSION_STRINGS", "pooled-1 pooled-2")
+    monkeypatch.delenv("TELEGRAM_SESSION_STRING", raising=False)
+    monkeypatch.delenv("TELEGRAM_SESSION_NAME", raising=False)
+    monkeypatch.setattr(runtime, "_acquire_session", lambda pool: pool[0])
+    monkeypatch.setattr(runtime, "StringSession", lambda value=None: f"str::{value}")
+    captured = {}
+
+    def _fake_build_client(session, label):
+        captured[label] = session
+        return object()
+
+    monkeypatch.setattr(runtime, "_build_client", _fake_build_client)
+
+    accounts = runtime._discover_accounts()
+
+    assert "default" in accounts
+    assert captured["default"] == "str::pooled-1"
+
+
+# --- runner retries a transient AuthKeyDuplicatedError -----------------------
+
+
+class _FlakyClient:
+    def __init__(self, fail_times):
+        self.fail_times = fail_times
+        self.connects = 0
+        self.disconnects = 0
+
+    async def connect(self):
+        self.connects += 1
+        if self.connects <= self.fail_times:
+            from telethon.errors import AuthKeyDuplicatedError
+
+            raise AuthKeyDuplicatedError(request=None)
+
+    async def disconnect(self):
+        self.disconnects += 1
+
+    async def is_user_authorized(self):
+        return True
+
+
+@pytest.fixture
+def no_sleep(monkeypatch):
+    async def _noop(_delay):
+        return None
+
+    monkeypatch.setattr(runner.asyncio, "sleep", _noop)
+
+
+@pytest.mark.asyncio
+async def test_connect_recovers_after_transient_authkey_duplicated(no_sleep):
+    client = _FlakyClient(fail_times=2)
+
+    await runner._connect_authorized_client("default", client)
+
+    assert client.connects == 3
+    assert client.disconnects == 2
+
+
+@pytest.mark.asyncio
+async def test_connect_reraises_authkey_duplicated_after_max_attempts(no_sleep):
+    from telethon.errors import AuthKeyDuplicatedError
+
+    client = _FlakyClient(fail_times=99)
+
+    with pytest.raises(AuthKeyDuplicatedError):
+        await runner._connect_authorized_client("default", client)
+
+    assert client.connects == 4