diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/redis_distributed_lock.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/redis_distributed_lock.py new file mode 100644 index 00000000000..37ad6a00282 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/redis_distributed_lock.py @@ -0,0 +1,62 @@ +"""Concrete ``DistributedLock`` over a Redis client: ``SET NX PX`` / ``DEL`` / ``EXISTS``. + +The cross-replica lock the ``RedisRefreshCoordinator`` elects refreshers with. ``acquire`` is an +atomic ``SET key NX PX ttl`` (only the first caller wins; the entry self-expires so a crashed holder +can't wedge refresh), ``release`` is ``DEL``, ``is_held`` is ``EXISTS``. The Redis client is injected +(in production the async client from LiteLLM's ``RedisCache``), so the lock is unit-testable with a +fake. A transport error on ``acquire`` is treated as "not acquired" so a Redis blip degrades to an +extra refresh, never a wedged one. +""" + +from __future__ import annotations + +from typing import Protocol + +from litellm._logging import verbose_logger + + +class RedisCommands(Protocol): + """The slice of the async Redis client this lock needs.""" + + async def set( + self, name: str, value: str, *, nx: bool = False, px: int | None = None + ) -> object | None: ... + + async def delete(self, *names: str) -> int: ... + + async def exists(self, *names: str) -> int: ... + + +class RedisDistributedLock: + def __init__(self, client: RedisCommands) -> None: + self._client = client + + async def acquire(self, key: str, ttl_seconds: float) -> bool: + try: + result = await self._client.set( + key, "1", nx=True, px=int(ttl_seconds * 1000) + ) + # Degrade on any Redis client error: redis.exceptions narrows only via an import that + # is Unknown under basedpyright, and the lock must never crash the resolve path. + except Exception as exc: # noqa: BLE001 + verbose_logger.warning("RedisDistributedLock.acquire failed: %s", exc) + return False + return result is not None + + async def release(self, key: str) -> None: + try: + await self._client.delete(key) + # Degrade on any Redis client error: redis.exceptions narrows only via an import that + # is Unknown under basedpyright, and the lock must never crash the resolve path. + except Exception as exc: # noqa: BLE001 + verbose_logger.warning("RedisDistributedLock.release failed: %s", exc) + + async def is_held(self, key: str) -> bool: + try: + return await self._client.exists(key) > 0 + # Degrade on any Redis client error: redis.exceptions narrows only via an import that + # is Unknown under basedpyright, and the lock must never crash the resolve path. + except Exception as exc: # noqa: BLE001 + # On error, report "not held" so a waiter stops waiting and re-reads rather than blocking. + verbose_logger.warning("RedisDistributedLock.is_held failed: %s", exc) + return False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_redis_distributed_lock.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_redis_distributed_lock.py new file mode 100644 index 00000000000..c91eccaf58a --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_redis_distributed_lock.py @@ -0,0 +1,79 @@ +"""Tests for the Redis SET NX PX lock: acquire semantics, release, is_held, error degradation.""" + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import ( + RedisDistributedLock, +) + + +class _FakeRedis: + def __init__(self, set_returns=True, exists_returns=1, raise_on=()): + self._set_returns = set_returns + self._exists_returns = exists_returns + self._raise_on = set(raise_on) + self.set_calls = [] + self.deleted = [] + + async def set(self, name, value, *, nx=False, px=None): + if "set" in self._raise_on: + raise RuntimeError("redis down") + self.set_calls.append((name, value, nx, px)) + return self._set_returns + + async def delete(self, *names): + self.deleted.extend(names) + return len(names) + + async def exists(self, *names): + if "exists" in self._raise_on: + raise RuntimeError("redis down") + return self._exists_returns + + +@pytest.mark.asyncio +async def test_acquire_uses_set_nx_px_and_reports_success(): + redis = _FakeRedis(set_returns=True) + lock = RedisDistributedLock(redis) + assert await lock.acquire("k", 10.0) is True + assert redis.set_calls == [("k", "1", True, 10000)] # NX + px in milliseconds + + +@pytest.mark.asyncio +async def test_acquire_reports_failure_when_key_already_held(): + # redis SET NX returns None when the key exists -> not acquired. + assert ( + await RedisDistributedLock(_FakeRedis(set_returns=None)).acquire("k", 10.0) + is False + ) + + +@pytest.mark.asyncio +async def test_acquire_degrades_to_not_acquired_on_redis_error(): + assert ( + await RedisDistributedLock(_FakeRedis(raise_on=["set"])).acquire("k", 10.0) + is False + ) + + +@pytest.mark.asyncio +async def test_release_deletes_the_key(): + redis = _FakeRedis() + await RedisDistributedLock(redis).release("k") + assert redis.deleted == ["k"] + + +@pytest.mark.asyncio +async def test_is_held_reflects_exists(): + assert await RedisDistributedLock(_FakeRedis(exists_returns=1)).is_held("k") is True + assert ( + await RedisDistributedLock(_FakeRedis(exists_returns=0)).is_held("k") is False + ) + + +@pytest.mark.asyncio +async def test_is_held_degrades_to_false_on_redis_error(): + assert ( + await RedisDistributedLock(_FakeRedis(raise_on=["exists"])).is_held("k") + is False + )