Skip to content

Commit 02f93bf

Browse files
authored
Merge pull request chigwell#215 from mokhtarabadi/fix/single-host-shared-lock
feat(singleton): shared session lock mode (rebase of chigwell#209 onto main, live-verified)
2 parents c9460f8 + 631699b commit 02f93bf

5 files changed

Lines changed: 205 additions & 17 deletions

File tree

‎.env.example‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,16 @@ TELEGRAM_SESSION_NAME=telegram_session
2929
# TELEGRAM_SESSION_STRING for the default account.
3030
# TELEGRAM_SESSION_STRINGS=<session A> <session B> <session C>
3131

32+
# --- Session lock (optional) ---
33+
# Before connecting, each server process takes a per-session lock so a second
34+
# instance (e.g. a connector restart) never connects the same session while the
35+
# first is live. exclusive (default): refuse to start if another instance still
36+
# holds the session after TELEGRAM_LOCK_GRACE_SECONDS (default 20). shared: let
37+
# several instances on this host use one session -- safe only when they all
38+
# reach Telegram from the same IP; if unsure, use the pool above instead.
39+
# TELEGRAM_SESSION_LOCK=exclusive
40+
# TELEGRAM_LOCK_GRACE_SECONDS=20
41+
3242
# --- Device identity (optional) ---
3343
# Controls how this client appears in Telegram > Settings > Devices. If unset,
3444
# Telethon falls back to the host platform (e.g. "arm64"). Set these so the

‎README.md‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,25 @@ rather than reusing a session another client holds — reuse would make Telegram
379379
permanently invalidate that session for both clients.
380380
- "Send this from my work account to @example"
381381

382+
### Sharing one session from one host
383+
384+
Before connecting, every server process takes a per-session lock (see the
385+
`AuthKeyDuplicatedError` entry under Troubleshooting). By default it is
386+
exclusive: a second instance for the same session waits briefly for the first
387+
to exit and otherwise refuses to start. That lock can only see processes on the
388+
same host, and Telegram's rule is about IPs, not processes: instances on one
389+
host only collide when they reach Telegram from different IPs (dual-stack or
390+
split-tunnel VPN hosts). When they all share one egress IP — a laptop running
391+
several MCP clients — you can let them share the session instead of pooling:
392+
393+
```env
394+
TELEGRAM_SESSION_LOCK=shared
395+
```
396+
397+
Shared instances coexist with each other but never overlap an exclusive one
398+
(except on Windows, where a shared instance takes no lock). If you are not sure
399+
every client egresses from the same IP, use the pool.
400+
382401
## Device Identity
383402

384403
These optional variables control how the client appears in Telegram under
@@ -645,7 +664,7 @@ Telegram messages, display names, chat titles, and button labels are untrusted c
645664
interactive phone-code login over stdio.
646665
- **Invalid API credentials:** verify `TELEGRAM_API_ID` and `TELEGRAM_API_HASH` at [my.telegram.org/apps](https://my.telegram.org/apps).
647666
- **Database is locked:** prefer string sessions, or make sure no other process is using the same file session.
648-
- **`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.
667+
- **`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).
649668
- **File tools are disabled:** pass allowed roots or configure MCP Roots in your client.
650669
- **Path rejected:** ensure the path is inside an allowed root and does not use traversal or wildcard patterns.
651670
- **Auth errors after password changes:** regenerate your session string.

‎telegram_mcp/runner.py‎

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,37 @@ def _lock_grace_seconds() -> float:
3535
return DEFAULT_GRACE_SECONDS
3636

3737

38+
_SESSION_LOCK_MODES = ("exclusive", "shared")
39+
40+
41+
def _session_lock_shared() -> bool:
42+
"""``TELEGRAM_SESSION_LOCK``: exclusive (default) | shared.
43+
44+
``exclusive`` refuses to start while another instance on this host holds
45+
the same session. ``shared`` takes the lock in shared mode instead, so any
46+
number of instances on this host can use one session -- safe only when
47+
they all reach Telegram from the same IP, since Telegram invalidates a
48+
session it sees from two IPs at once. Shared and exclusive instances never
49+
overlap, so the default keeps its guarantee.
50+
"""
51+
raw = (os.getenv("TELEGRAM_SESSION_LOCK") or "exclusive").strip().lower()
52+
if raw not in _SESSION_LOCK_MODES:
53+
accepted = ", ".join(_SESSION_LOCK_MODES)
54+
raise SystemExit(f"Invalid TELEGRAM_SESSION_LOCK '{raw}'. Expected one of: {accepted}.")
55+
return raw == "shared"
56+
57+
3858
async def _connect_authorized_client(label, client) -> None:
39-
# First, prevent our own duplicate-spawn case outright: an exclusive
40-
# per-session lock means a second instance of this server never even
41-
# attempts to connect while another instance already holds the same
42-
# session (see telegram_mcp/singleton.py for why and how).
59+
# First, prevent our own duplicate-spawn case outright: a per-session lock
60+
# means a second instance of this server never even attempts to connect
61+
# while another instance already holds the same session (see
62+
# telegram_mcp/singleton.py for why and how). The lock only sees processes
63+
# on this host; TELEGRAM_SESSION_LOCK=shared lets those share one session
64+
# where they all reach Telegram from a single IP.
4365
lock = SessionLock(label, session_identity(client))
44-
await asyncio.to_thread(lock.acquire, grace_seconds=_lock_grace_seconds())
66+
await asyncio.to_thread(
67+
lock.acquire, grace_seconds=_lock_grace_seconds(), shared=_session_lock_shared()
68+
)
4569
_session_locks[label] = lock
4670

4771
# Once we hold the lock, still tolerate a transient AuthKeyDuplicatedError
@@ -173,7 +197,10 @@ async def _warm_caches() -> None:
173197
"session (e.g. the client restarted the connector without the old "
174198
"process exiting yet). This instance is exiting instead of "
175199
"connecting a second time, which would risk Telegram invalidating "
176-
"the session for both. Retry once the other instance is gone.",
200+
"the session for both. Retry once the other instance is gone, or "
201+
"set TELEGRAM_SESSION_LOCK=shared if several instances on this "
202+
"host are meant to share one session (safe when they all reach "
203+
"Telegram from the same IP).",
177204
file=sys.stderr,
178205
)
179206
sys.exit(1)
@@ -193,6 +220,7 @@ def main() -> None:
193220
_configure_allowed_roots_from_cli(sys.argv[1:])
194221
_runtime._apply_exposed_tools_mode()
195222
_transcription.validate_transcription_config()
223+
_session_lock_shared() # fail loudly at startup on a bad toggle
196224
asyncio.run(_main())
197225

198226

‎telegram_mcp/singleton.py‎

Lines changed: 64 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,16 @@
2020
"restart replaces the old process" case) and, if the lock is still held once
2121
the grace period elapses, raises :class:`SessionLockError` instead of ever
2222
calling ``connect()`` -- so the live session is never disturbed.
23+
24+
The lock can only ever see processes on this host, and those only collide
25+
at Telegram when they reach it from different IPs (dual-stack or
26+
split-tunnel hosts). Where every instance shares one egress IP,
27+
``acquire(shared=True)`` takes the lock in shared mode instead: shared
28+
holders coexist with each other, but a shared holder and an exclusive
29+
holder never overlap, so an operator who never opted in keeps the
30+
exclusive guarantee. (On Windows, ``msvcrt`` byte-range locks have no
31+
shared mode; a shared holder there takes no lock at all.) The exclusive
32+
holder leaves its PID in the lock file so a refused contender can name it.
2333
"""
2434

2535
from __future__ import annotations
@@ -34,7 +44,9 @@
3444
if os.name == "nt":
3545
import msvcrt
3646

37-
def _try_lock(fh: IO) -> bool:
47+
def _try_lock(fh: IO, shared: bool = False) -> bool:
48+
if shared:
49+
return True # no shared byte-range locks on Windows; see module docstring
3850
try:
3951
fh.seek(0)
4052
msvcrt.locking(fh.fileno(), msvcrt.LK_NBLCK, 1)
@@ -52,9 +64,10 @@ def _unlock(fh: IO) -> None:
5264
else:
5365
import fcntl
5466

55-
def _try_lock(fh: IO) -> bool:
67+
def _try_lock(fh: IO, shared: bool = False) -> bool:
5668
try:
57-
fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
69+
mode = fcntl.LOCK_SH if shared else fcntl.LOCK_EX
70+
fcntl.flock(fh.fileno(), mode | fcntl.LOCK_NB)
5871
return True
5972
except OSError:
6073
return False
@@ -97,35 +110,76 @@ def acquire(
97110
*,
98111
grace_seconds: float = DEFAULT_GRACE_SECONDS,
99112
poll_interval: float = DEFAULT_POLL_INTERVAL,
113+
shared: bool = False,
100114
) -> None:
101115
"""Block (up to ``grace_seconds``) until the lock is free, then take it.
102116
103-
Raises :class:`SessionLockError` if another live process still holds
104-
the lock once the grace period elapses.
117+
``shared`` takes the lock in shared mode: any number of shared holders
118+
coexist, but a shared holder and an exclusive holder never overlap.
119+
Raises :class:`SessionLockError` if the lock is still held once the
120+
grace period elapses.
105121
"""
106122
fh = open(self.path, "a+")
107123
deadline = time.monotonic() + grace_seconds
108124
while True:
109-
if _try_lock(fh):
125+
if _try_lock(fh, shared=shared):
110126
self._fh = fh
127+
self._record_holder(shared)
111128
return
112129
if time.monotonic() >= deadline:
113130
fh.close()
131+
holder = self.holder_pid()
132+
where = (
133+
f"lock {self.path} held by PID {holder}"
134+
if holder
135+
else f"lock held: {self.path}"
136+
)
114137
raise SessionLockError(
115138
"Another telegram-mcp process is already connected with this "
116-
f"session (lock held: {self.path}). Refusing to connect a "
117-
"second time to avoid Telegram's AuthKeyDuplicatedError. If "
118-
"that other process already exited, this lock will clear on "
119-
"its own -- retry."
139+
f"session ({where}). Refusing to connect a second time to "
140+
"avoid Telegram's AuthKeyDuplicatedError. If that other "
141+
"process already exited, this lock will clear on its own -- "
142+
"retry."
120143
)
121144
time.sleep(poll_interval)
122145

123146
def release(self) -> None:
124147
if self._fh is not None:
148+
self._clear_holder()
125149
_unlock(self._fh)
126150
self._fh.close()
127151
self._fh = None
128152

153+
def holder_pid(self) -> Optional[int]:
154+
"""PID left in the lock file by the current exclusive holder, if any.
155+
156+
Best effort: Windows refuses to read a locked byte range, and a shared
157+
holder records nothing (there may be several).
158+
"""
159+
try:
160+
text = self.path.read_text().strip()
161+
except OSError:
162+
return None
163+
return int(text) if text.isdigit() else None
164+
165+
def _record_holder(self, shared: bool) -> None:
166+
# Exclusive holders leave their PID for a refused contender to report;
167+
# shared holders wipe whatever a previous exclusive holder left behind.
168+
self._clear_holder()
169+
if not shared:
170+
try:
171+
self._fh.write(str(os.getpid()))
172+
self._fh.flush()
173+
except OSError:
174+
pass
175+
176+
def _clear_holder(self) -> None:
177+
try:
178+
self._fh.seek(0)
179+
self._fh.truncate()
180+
except OSError:
181+
pass
182+
129183

130184
def session_identity(client: object) -> str:
131185
"""Derive a stable identity string for a ``TelegramClient``'s session.

‎tests/test_runner.py‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import os
2+
13
import pytest
24

35
from telegram_mcp import runner
@@ -43,6 +45,9 @@ def _init_with_tmp_dir(self, label, session_identity, *, lock_dir=tmp_path):
4345

4446
monkeypatch.setattr(singleton_module.SessionLock, "__init__", _init_with_tmp_dir)
4547
monkeypatch.setattr(runner, "_lock_grace_seconds", lambda: 0.01)
48+
# load_dotenv() may have pulled TELEGRAM_SESSION_LOCK from the developer's
49+
# .env; tests assume the exclusive default unless they set it themselves.
50+
monkeypatch.delenv("TELEGRAM_SESSION_LOCK", raising=False)
4651
yield
4752
runner._session_locks.clear()
4853

@@ -96,6 +101,78 @@ async def test_connect_authorized_client_allows_different_sessions_concurrently(
96101
assert second.connected is True
97102

98103

104+
@pytest.mark.asyncio
105+
async def test_shared_lock_mode_lets_instances_share_a_session(monkeypatch):
106+
monkeypatch.setenv("TELEGRAM_SESSION_LOCK", "shared")
107+
first = _FakeClient(authorized=True, identity="shared-session")
108+
second = _FakeClient(authorized=True, identity="shared-session")
109+
110+
await runner._connect_authorized_client("default", first)
111+
first_lock = runner._session_locks["default"]
112+
await runner._connect_authorized_client("default", second)
113+
114+
assert first.connected is True
115+
assert second.connected is True
116+
117+
first_lock.release()
118+
runner._session_locks["default"].release()
119+
120+
121+
@pytest.mark.asyncio
122+
@pytest.mark.parametrize(
123+
"first_mode, second_mode", [("exclusive", "shared"), ("shared", "exclusive")]
124+
)
125+
async def test_shared_and_exclusive_instances_never_overlap(monkeypatch, first_mode, second_mode):
126+
first = _FakeClient(authorized=True, identity="shared-session")
127+
second = _FakeClient(authorized=True, identity="shared-session")
128+
129+
monkeypatch.setenv("TELEGRAM_SESSION_LOCK", first_mode)
130+
await runner._connect_authorized_client("default", first)
131+
first_lock = runner._session_locks["default"]
132+
133+
monkeypatch.setenv("TELEGRAM_SESSION_LOCK", second_mode)
134+
with pytest.raises(runner.SessionLockError, match="already connected"):
135+
await runner._connect_authorized_client("default", second)
136+
137+
assert second.connected is False
138+
first_lock.release()
139+
140+
141+
@pytest.mark.asyncio
142+
async def test_lock_error_names_the_exclusive_holder():
143+
first = _FakeClient(authorized=True, identity="shared-session")
144+
second = _FakeClient(authorized=True, identity="shared-session")
145+
146+
await runner._connect_authorized_client("default", first)
147+
lock = runner._session_locks["default"]
148+
149+
with pytest.raises(runner.SessionLockError, match=f"held by PID {os.getpid()}"):
150+
await runner._connect_authorized_client("default", second)
151+
152+
lock.release()
153+
assert lock.path.read_text() == "" # a released lock names nobody
154+
155+
156+
@pytest.mark.parametrize(
157+
"value, shared",
158+
[(None, False), ("", False), ("exclusive", False), ("shared", True), (" Shared ", True)],
159+
)
160+
def test_session_lock_mode_parsing(monkeypatch, value, shared):
161+
if value is None:
162+
monkeypatch.delenv("TELEGRAM_SESSION_LOCK", raising=False)
163+
else:
164+
monkeypatch.setenv("TELEGRAM_SESSION_LOCK", value)
165+
166+
assert runner._session_lock_shared() is shared
167+
168+
169+
def test_session_lock_mode_rejects_unknown_values(monkeypatch):
170+
monkeypatch.setenv("TELEGRAM_SESSION_LOCK", "sometimes")
171+
172+
with pytest.raises(SystemExit, match="Invalid TELEGRAM_SESSION_LOCK 'sometimes'"):
173+
runner._session_lock_shared()
174+
175+
99176
class _FakeSettings:
100177
def __init__(self):
101178
self.host = None

0 commit comments

Comments
 (0)