Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
20 changes: 20 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
31 changes: 30 additions & 1 deletion telegram_mcp/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
102 changes: 97 additions & 5 deletions telegram_mcp/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"

Expand All @@ -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(
Expand Down
144 changes: 144 additions & 0 deletions tests/test_session_pool.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,144 @@
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
Loading