Skip to content
Merged
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
16 changes: 9 additions & 7 deletions telegram_mcp/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@
except ImportError: # pragma: no cover - Windows fallback
fcntl = None

from telegram_mcp.singleton import try_lock_exclusive

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 @@ -366,11 +368,6 @@ def _parse_session_pool() -> List[str]:

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)
Expand All @@ -380,9 +377,12 @@ def _acquire_session(pool: List[str]) -> str:
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)
# "a+", not "w": on Windows the lock covers the first byte, and
# truncating a file another live client holds is refused.
fh = open(lock_path, "a+")
except OSError:
continue
if not try_lock_exclusive(fh):
# Locked by another live client — try the next session.
try:
fh.close()
Expand All @@ -391,6 +391,8 @@ def _acquire_session(pool: List[str]) -> str:
continue
_SESSION_LOCKS.append(fh)
try:
fh.seek(0)
fh.truncate()
fh.write(f"pid={os.getpid()}\n")
fh.flush()
except OSError:
Expand Down
8 changes: 8 additions & 0 deletions telegram_mcp/singleton.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,14 @@ def _unlock(fh: IO) -> None:
pass


# Public names for the primitives above. The session pool in ``runtime`` needs
# the same "try to take an exclusive lock, don't block" behaviour on both
# platforms, and duplicating the msvcrt/fcntl split there is how one of the two
# copies ends up POSIX-only.
try_lock_exclusive = _try_lock
release_lock = _unlock


DEFAULT_LOCK_DIR = Path(tempfile.gettempdir()) / "telegram-mcp-locks"
DEFAULT_GRACE_SECONDS = 20.0
DEFAULT_POLL_INTERVAL = 0.5
Expand Down
21 changes: 15 additions & 6 deletions tests/test_session_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import pytest

from telegram_mcp import runner, runtime
from telegram_mcp.singleton import try_lock_exclusive

# --- _parse_session_pool -----------------------------------------------------

Expand All @@ -30,13 +31,11 @@ def isolated_lock_dir(tmp_path, monkeypatch):

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)
fh = open(path, "a+")
assert try_lock_exclusive(fh), "test could not take the lock it means to hold"
return fh


Expand Down Expand Up @@ -64,9 +63,19 @@ def test_acquire_session_raises_when_pool_exhausted(isolated_lock_dir):
fh.close()


def test_acquire_session_without_fcntl_uses_first(isolated_lock_dir, monkeypatch):
monkeypatch.setattr(runtime, "fcntl", None)
def test_acquire_session_locks_are_visible_to_other_clients(isolated_lock_dir):
"""The claimed slot must actually be locked, on every platform.

Windows previously had no advisory locking here and every client was handed
pool[0], which is precisely the collision the pool exists to prevent.
"""
assert runtime._acquire_session(["AAA", "BBB"]) == "AAA"
digest = hashlib.sha1(b"AAA").hexdigest()[:16]
path = os.path.join(
str(isolated_lock_dir), "telegram-mcp-session-locks", f"session-{digest}.lock"
)
with open(path, "a+") as rival:
assert not try_lock_exclusive(rival)


# --- _discover_accounts prefers the pool for the default account -------------
Expand Down
Loading