Merge pull request #186 from diyor28/fix/bound-reconnect-and-pool-exhaustion

fix: bound reconnect and fail fast when the session pool is exhausted

Author: chigwell

README.md

@@ -329,6 +329,11 @@ 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.
+
+Size the pool to the number of clients you actually run concurrently. If every
+slot is already claimed, the server refuses to start with an explicit error
+rather than reusing a session another client holds — reuse would make Telegram
+permanently invalidate that session for both clients.
 - "Send this from my work account to @example"
 
 ## Device Identity

telegram_mcp/runtime.py

@@ -23,6 +23,7 @@
 from mcp.shared.exceptions import McpError
 from pythonjsonlogger import jsonlogger
 from telethon import TelegramClient, functions, types, utils
+from telethon.errors import AuthKeyDuplicatedError
 from telethon.sessions import StringSession
 from telethon.tl.types import (
     User,
@@ -396,13 +397,16 @@ def _acquire_session(pool: List[str]) -> str:
             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,
+    # Handing out an already-claimed session here would make Telegram burn it
+    # with AuthKeyDuplicatedError — losing the slot for the client that owns it
+    # too. Refusing to start is recoverable; a burned session is not.
+    raise RuntimeError(
+        f"All {len(pool)} pooled Telegram session(s) are already claimed by other "
+        "live clients, so this one has no session to use. Add another session to "
+        "TELEGRAM_SESSION_STRINGS (generate it with "
+        "`uv run session_string_generator.py`) — one slot per concurrent client — "
+        "or stop one of the other clients."
     )
-    return pool[0]
 
 
 def _discover_accounts() -> dict[str, TelegramClient]:
@@ -524,6 +528,7 @@ async def _call_for(label):
 
 _last_conn_verified: dict[int, float] = {}
 _CONN_VERIFY_INTERVAL: float = 30.0  # seconds between live pings
+_RECONNECT_TIMEOUT: float = 30.0  # seconds before a reconnect attempt is abandoned
 
 
 async def _force_reconnect(cl: TelegramClient):
@@ -534,10 +539,26 @@ async def _force_reconnect(cl: TelegramClient):
         await cl.disconnect()
     except Exception:
         pass
-    await cl.connect()
+    try:
+        await asyncio.wait_for(cl.connect(), timeout=_RECONNECT_TIMEOUT)
+    except AuthKeyDuplicatedError as exc:
+        # Telegram permanently invalidates an auth key used from two IPs at
+        # once, so retrying here can never succeed — surface it instead of
+        # letting the caller sit in a reconnect loop.
+        raise RuntimeError(
+            "Telegram session is no longer usable: the same session string was "
+            "used by another client at the same time (AuthKeyDuplicatedError). "
+            "Give each concurrent client its own session via "
+            "TELEGRAM_SESSION_STRINGS or TELEGRAM_SESSION_STRING_<LABEL>, then "
+            "regenerate the burned session with `uv run session_string_generator.py`."
+        ) from exc
+    except asyncio.TimeoutError as exc:
+        raise RuntimeError(
+            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 cl.start()
+        await asyncio.wait_for(cl.start(), timeout=_RECONNECT_TIMEOUT)
     _last_conn_verified[id(cl)] = time.time()
     reconnect_logger.warning("Forced reconnect successful")
 

tests/test_runtime.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from datetime import datetime, timezone
 from pathlib import Path
@@ -430,6 +431,37 @@ async def test_ensure_connected_skips_recently_verified_client(monkeypatch):
     assert client.calls == ["is_connected"]
 
 
+class _HangingConnectClient(_ConnectivityClient):
+    async def connect(self):
+        self.calls.append("connect")
+        await asyncio.sleep(3600)
+
+
+class _DuplicatedKeyClient(_ConnectivityClient):
+    async def connect(self):
+        from telethon.errors import AuthKeyDuplicatedError
+
+        self.calls.append("connect")
+        raise AuthKeyDuplicatedError(request=None)
+
+
+@pytest.mark.asyncio
+async def test_force_reconnect_times_out_instead_of_hanging(monkeypatch):
+    client = _HangingConnectClient(connected=False, authorized=True)
+    monkeypatch.setattr(runtime, "_RECONNECT_TIMEOUT", 0.01)
+
+    with pytest.raises(RuntimeError, match="timed out"):
+        await runtime._force_reconnect(client)
+
+
+@pytest.mark.asyncio
+async def test_force_reconnect_reports_burned_session(monkeypatch):
+    client = _DuplicatedKeyClient(connected=False, authorized=True)
+
+    with pytest.raises(RuntimeError, match="no longer usable"):
+        await runtime._force_reconnect(client)
+
+
 class _ResolvingClient:
     def __init__(self, method_name, failures):
         self.method_name = method_name

tests/test_session_pool.py

@@ -52,11 +52,13 @@ def test_acquire_session_skips_slot_locked_by_another_client(isolated_lock_dir):
         foreign.close()
 
 
-def test_acquire_session_falls_back_to_first_when_pool_exhausted(isolated_lock_dir):
+def test_acquire_session_raises_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"
+        # Only two slots exist and both are taken -> refuse rather than hand out
+        # a session another live client is already using.
+        with pytest.raises(RuntimeError, match="already claimed"):
+            runtime._acquire_session(["AAA", "BBB"])
     finally:
         for fh in held:
             fh.close()