From 0b0fcd06de1ddd81263d42713e65004c5a418a58 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 19:40:36 +0000 Subject: [PATCH] fix(redis): send MCP marker and refresh lock scripts with plain EVAL --- .../redis/stores/mcp_credential_refresh.py | 4 +-- litellm/redis/stores/mcp_single_use.py | 2 +- .../proxy/mcp_server/test_gateway_dcr_flow.py | 4 +-- .../stores/test_mcp_credential_refresh.py | 34 +++++++++++++++++++ .../unit/redis/stores/test_mcp_single_use.py | 26 ++++++++++++++ 5 files changed, 65 insertions(+), 5 deletions(-) create mode 100644 tests/unit/redis/stores/test_mcp_credential_refresh.py create mode 100644 tests/unit/redis/stores/test_mcp_single_use.py diff --git a/litellm/redis/stores/mcp_credential_refresh.py b/litellm/redis/stores/mcp_credential_refresh.py index 2950ab7f5d3..43f112fd049 100644 --- a/litellm/redis/stores/mcp_credential_refresh.py +++ b/litellm/redis/stores/mcp_credential_refresh.py @@ -40,11 +40,11 @@ class TransportMcpRefreshLockStore: return await self.transport.set(_lock_key(target), owner, ttl=ttl, nx=True) async def extend(self, target: RefreshLockTarget, owner: str, *, ttl: timedelta) -> bool | RedisFailure: - reply: Final = await self.transport.eval(_EXTEND_IF_OWNER, (_lock_key(target),), (owner, ttl_ms(ttl))) + reply: Final = await self.transport.eval_direct(_EXTEND_IF_OWNER, (_lock_key(target),), (owner, ttl_ms(ttl))) return reply if isinstance(reply, RedisFailure) else reply == 1 async def release(self, target: RefreshLockTarget, owner: str) -> bool | RedisFailure: - reply: Final = await self.transport.eval(_RELEASE_IF_OWNER, (_lock_key(target),), (owner,)) + reply: Final = await self.transport.eval_direct(_RELEASE_IF_OWNER, (_lock_key(target),), (owner,)) return reply if isinstance(reply, RedisFailure) else reply == 1 async def is_held(self, target: RefreshLockTarget) -> bool | RedisFailure: diff --git a/litellm/redis/stores/mcp_single_use.py b/litellm/redis/stores/mcp_single_use.py index 9a5d6643369..500af7a46b4 100644 --- a/litellm/redis/stores/mcp_single_use.py +++ b/litellm/redis/stores/mcp_single_use.py @@ -74,7 +74,7 @@ class TransportMcpSingleUseStore: transport: RedisTransport async def claim(self, marker: SingleUseMarker, *, ttl: timedelta) -> int | RedisFailure: - reply: Final = await self.transport.eval(_CLAIM, (marker_key(marker),), (1, ttl // timedelta(seconds=1))) + reply: Final = await self.transport.eval_direct(_CLAIM, (marker_key(marker),), (1, ttl // timedelta(seconds=1))) return reply if isinstance(reply, RedisFailure) else _count(reply) async def peek(self, marker: SingleUseMarker) -> bool | RedisFailure: diff --git a/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py index 8228d9f2c5f..6ef4244035f 100644 --- a/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py @@ -2473,11 +2473,11 @@ class _RedisThatFaultsOnTheSecondClaim(fakeredis.FakeAsyncRedis): super().__init__(server=server) self.claims = 0 - async def evalsha(self, sha: str, numkeys: int, *keys_and_args: str | bytes | float) -> object: + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> object: self.claims += 1 if self.claims == 2: raise RedisConnectionError("redis down") - return await super().evalsha(sha, numkeys, *keys_and_args) + return await super().eval(script, numkeys, *keys_and_args) @pytest.mark.asyncio diff --git a/tests/unit/redis/stores/test_mcp_credential_refresh.py b/tests/unit/redis/stores/test_mcp_credential_refresh.py new file mode 100644 index 00000000000..4f27c52a45a --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_credential_refresh.py @@ -0,0 +1,34 @@ +from datetime import timedelta +from typing import Final + +import fakeredis +import pytest +from redis.exceptions import ResponseError + +from litellm.redis.stores.mcp_credential_refresh import RefreshLockTarget, TransportMcpRefreshLockStore +from tests._support.fake_redis import fake_redis_transport + +TARGET: Final = RefreshLockTarget(user_id="u", server_id="s") + + +class _RedisWithoutEvalsha(fakeredis.FakeAsyncRedis): + async def evalsha(self, sha: str, numkeys: int, *keys_and_args: str | bytes | float) -> object: + raise ResponseError("NOPERM this user has no permissions to run the 'evalsha' command") + + +@pytest.mark.asyncio +async def test_extend_and_release_honor_the_owner_with_plain_eval_when_evalsha_is_denied() -> None: + raw: Final = _RedisWithoutEvalsha(server=fakeredis.FakeServer()) + store: Final = TransportMcpRefreshLockStore(fake_redis_transport(raw, failure_threshold=1, namespace="ns")) + assert await store.acquire(TARGET, "owner-a", ttl=timedelta(seconds=30)) is True + lock_key: Final = "ns:mcp:refresh_lock:u:s" + + assert await store.extend(TARGET, "owner-b", ttl=timedelta(seconds=300)) is False + assert await store.release(TARGET, "owner-b") is False + assert await raw.get(lock_key) == b"owner-a" + assert await raw.pttl(lock_key) <= 30_000 + + assert await store.extend(TARGET, "owner-a", ttl=timedelta(seconds=300)) is True + assert 30_000 < await raw.pttl(lock_key) <= 300_000 + assert await store.release(TARGET, "owner-a") is True + assert await raw.get(lock_key) is None diff --git a/tests/unit/redis/stores/test_mcp_single_use.py b/tests/unit/redis/stores/test_mcp_single_use.py new file mode 100644 index 00000000000..2d799dbcdeb --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_single_use.py @@ -0,0 +1,26 @@ +from datetime import timedelta +from typing import Final + +import fakeredis +import pytest +from redis.exceptions import ResponseError + +from litellm.redis.stores.mcp_single_use import AuthCodeUsed, TransportMcpSingleUseStore +from tests._support.fake_redis import fake_redis_transport + + +class _RedisWithoutEvalsha(fakeredis.FakeAsyncRedis): + async def evalsha(self, sha: str, numkeys: int, *keys_and_args: str | bytes | float) -> object: + raise ResponseError("NOPERM this user has no permissions to run the 'evalsha' command") + + +@pytest.mark.asyncio +async def test_claim_counts_and_expires_with_plain_eval_when_evalsha_is_denied() -> None: + raw: Final = _RedisWithoutEvalsha(server=fakeredis.FakeServer()) + store: Final = TransportMcpSingleUseStore(fake_redis_transport(raw, failure_threshold=1, namespace="ns")) + marker: Final = AuthCodeUsed("jti-1") + + assert await store.claim(marker, ttl=timedelta(seconds=90)) == 1 + assert await store.claim(marker, ttl=timedelta(seconds=90)) == 2 + assert await raw.get("ns:mcp_gateway_dcr_code_used:jti-1") == b"2" + assert 0 < await raw.pttl("ns:mcp_gateway_dcr_code_used:jti-1") <= 90_000