From d73aff25a42d57b8e98cf68a098e7f6bdaaf4ca0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:35:36 +0000 Subject: [PATCH] test(redis): retain refresh lock compatibility coverage --- .../outbound_credentials/conftest.py | 53 ++++ .../test_redis_distributed_lock.py | 133 +++++++++ .../test_redis_refresh_coordinator.py | 256 +++++++++++++++++- 3 files changed, 439 insertions(+), 3 deletions(-) create mode 100644 tests/unit/proxy/mcp_server/outbound_credentials/conftest.py diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/conftest.py b/tests/unit/proxy/mcp_server/outbound_credentials/conftest.py new file mode 100644 index 00000000000..f0c2aa7734d --- /dev/null +++ b/tests/unit/proxy/mcp_server/outbound_credentials/conftest.py @@ -0,0 +1,53 @@ +from typing import Final, NoReturn + +import fakeredis +import pytest +import pytest_asyncio +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.mcp_server.outbound_credentials.redis_distributed_lock import RedisDistributedLock +from litellm.redis.stores.mcp_credential_refresh import TransportMcpRefreshLockStore +from tests._support.fake_redis import fake_redis_cache + + +def lost_write_reply(response: object, **kwargs: object) -> NoReturn: + raise RedisTimeoutError("reply lost after applying the write") + + +class RedisClock: + def __init__(self) -> None: + self.now: float = 1_700_000_000.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +@pytest.fixture +def redis_clock(monkeypatch: pytest.MonkeyPatch) -> RedisClock: + clock: Final = RedisClock() + monkeypatch.setattr("time.time", clock) + return clock + + +@pytest.fixture +def refresh_server(redis_clock: RedisClock) -> fakeredis.FakeServer: + return fakeredis.FakeServer() + + +@pytest.fixture +def refresh_raw(refresh_server: fakeredis.FakeServer) -> fakeredis.FakeAsyncRedis: + return fakeredis.FakeAsyncRedis(server=refresh_server) + + +@pytest_asyncio.fixture(loop_scope="function") +async def refresh_cache(refresh_server: fakeredis.FakeServer) -> RedisCache: + return fake_redis_cache(refresh_server, namespace="ns") + + +@pytest.fixture +def refresh_lock(refresh_cache: RedisCache) -> RedisDistributedLock: + return RedisDistributedLock(TransportMcpRefreshLockStore(refresh_cache.transport())) diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_distributed_lock.py b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_distributed_lock.py index 239c8d31d02..765257d1dee 100644 --- a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_distributed_lock.py +++ b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_distributed_lock.py @@ -1,8 +1,12 @@ """Tests for the Redis lock: acquire NX/PX with a token, compare-and-delete release, namespacing.""" +from typing import Final, Literal + import fakeredis import pytest +from litellm.caching.redis_cache import RedisCache +from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD from litellm.proxy.mcp_server.outbound_credentials.redis_distributed_lock import ( RedisDistributedLock, ) @@ -10,7 +14,9 @@ from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator imp LockAcquisition, ) from litellm.redis.stores.mcp_credential_refresh import RefreshLockTarget, TransportMcpRefreshLockStore +from litellm.redis.transport.contract import RedisFailure from tests._support.fake_redis import fake_redis_cache +from tests.unit.proxy.mcp_server.outbound_credentials.conftest import lost_write_reply TARGET = RefreshLockTarget(user_id="u", server_id="s") KEY = "mcp:refresh_lock:u:s" @@ -106,3 +112,130 @@ async def test_extend_release_and_is_held_degrade_on_redis_error(server: fakered server.connected = True assert await lock.is_held(TARGET) is True + + +_IDENTITY: Final = RefreshLockTarget("u", "s") +_KEY: Final = "ns:mcp:refresh_lock:u:s" +_OLD_RELEASE: Final = ( + "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("namespace", (None, "ns", "mcp")) +@pytest.mark.parametrize("owner", ('owner "雪" \\ 1.25', "", "123")) +async def test_acquire_preserves_raw_owner_and_fractional_ttl_for_old_readers( + refresh_server: fakeredis.FakeServer, refresh_raw: fakeredis.FakeAsyncRedis, namespace: str | None, owner: str +) -> None: + lock: Final = RedisDistributedLock( + TransportMcpRefreshLockStore(fake_redis_cache(refresh_server, namespace=namespace).transport()) + ) + key: Final = "ns:mcp:refresh_lock:u:s" if namespace == "ns" else "mcp:refresh_lock:u:s" + assert await lock.acquire(_IDENTITY, owner, 1.2345) is LockAcquisition.ACQUIRED + assert await refresh_raw.keys("*") == [key.encode()] + assert await refresh_raw.get(key) == owner.encode() + assert await refresh_raw.pttl(key) == 1235 + assert await refresh_raw.eval(_OLD_RELEASE, 1, key, owner) == 1 + assert await lock.is_held(_IDENTITY) is False + + +@pytest.mark.asyncio +async def test_acquire_reports_held_without_overwriting_owner_or_expiry( + refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis +) -> None: + assert await refresh_lock.acquire(_IDENTITY, "tok-a", 1.2345) is LockAcquisition.ACQUIRED + assert await refresh_lock.acquire(_IDENTITY, "tok-b", 10.0) is LockAcquisition.HELD + assert await refresh_raw.get(_KEY) == b"tok-a" + assert await refresh_raw.pttl(_KEY) == 1235 + + +@pytest.mark.asyncio +async def test_each_user_server_identity_has_an_independent_lease( + refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis +) -> None: + assert await refresh_lock.acquire(_IDENTITY, "first", 10.0) is LockAcquisition.ACQUIRED + assert await refresh_lock.acquire(RefreshLockTarget("u2", "s"), "second", 10.0) is LockAcquisition.ACQUIRED + assert await refresh_lock.acquire(RefreshLockTarget("u", "s2"), "third", 10.0) is LockAcquisition.ACQUIRED + await refresh_lock.release(_IDENTITY, "first") + assert await refresh_raw.mget(_KEY, "ns:mcp:refresh_lock:u2:s", "ns:mcp:refresh_lock:u:s2") == [ + None, + b"second", + b"third", + ] + + +@pytest.mark.asyncio +async def test_old_writer_lease_is_held_renewed_and_released_only_by_its_owner( + refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis +) -> None: + owner: Final = 'old "雪" \\ 1.25' + await refresh_raw.set(_KEY, owner, nx=True, px=1235) + assert await refresh_lock.is_held(_IDENTITY) is True + assert await refresh_lock.acquire(_IDENTITY, "new", 10.0) is LockAcquisition.HELD + assert await refresh_lock.extend(_IDENTITY, "new", 30.0) is False + await refresh_lock.release(_IDENTITY, "new") + assert await refresh_raw.get(_KEY) == owner.encode() + assert await refresh_raw.pttl(_KEY) == 1235 + assert await refresh_lock.extend(_IDENTITY, owner, 30.2345) is True + assert await refresh_raw.pttl(_KEY) == 30235 + await refresh_lock.release(_IDENTITY, owner) + assert await refresh_lock.is_held(_IDENTITY) is False + + +@pytest.mark.asyncio +async def test_outage_and_open_breaker_keep_acquisition_error_distinct_from_held( + refresh_server: fakeredis.FakeServer, + refresh_cache: RedisCache, + refresh_lock: RedisDistributedLock, + refresh_raw: fakeredis.FakeAsyncRedis, +) -> None: + refresh_server.connected = False + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + assert await refresh_lock.acquire(_IDENTITY, "tok", 10.0) is LockAcquisition.ERROR + refresh_server.connected = True + assert await refresh_lock.acquire(_IDENTITY, "tok", 10.0) is LockAcquisition.ERROR + failure: Final = await refresh_cache.transport().get("mcp:refresh_lock:u:s") + assert isinstance(failure, RedisFailure) and failure.kind == "breaker_open" + assert await refresh_raw.get(_KEY) is None + + +@pytest.mark.asyncio +async def test_extend_release_and_is_held_degrade_on_redis_error_with_redis_store( + refresh_server: fakeredis.FakeServer, refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis +) -> None: + assert await refresh_lock.acquire(_IDENTITY, "tok", 10.0) is LockAcquisition.ACQUIRED + refresh_server.connected = False + assert await refresh_lock.extend(_IDENTITY, "tok", 10.0) is False + await refresh_lock.release(_IDENTITY, "tok") + assert await refresh_lock.is_held(_IDENTITY) is False + refresh_server.connected = True + assert await refresh_raw.get(_KEY) == b"tok" + assert await refresh_lock.is_held(_IDENTITY) is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ("acquire", "extend", "release")) +async def test_lost_write_reply_does_not_claim_the_write_was_rejected( + refresh_cache: RedisCache, + refresh_lock: RedisDistributedLock, + refresh_raw: fakeredis.FakeAsyncRedis, + operation: Literal["acquire", "extend", "release"], +) -> None: + client: Final = refresh_cache.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + if operation == "acquire": + client.set_response_callback("SET", lost_write_reply) + assert await refresh_lock.acquire(_IDENTITY, "tok", 1.2345) is LockAcquisition.ERROR + assert await refresh_raw.get(_KEY) == b"tok" + assert await refresh_raw.pttl(_KEY) == 1235 + return + await refresh_raw.set(_KEY, "tok", px=1235) + client.set_response_callback("EVAL", lost_write_reply) + client.set_response_callback("EVALSHA", lost_write_reply) + if operation == "extend": + assert await refresh_lock.extend(_IDENTITY, "tok", 30.2345) is False + assert await refresh_raw.get(_KEY) == b"tok" + assert await refresh_raw.pttl(_KEY) == 30235 + return + await refresh_lock.release(_IDENTITY, "tok") + assert await refresh_raw.get(_KEY) is None diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py index 19e642df0c7..ce4dae8ac59 100644 --- a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py +++ b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py @@ -1,17 +1,22 @@ """Tests for the cross-replica refresh coordinator: winner refreshes, losers wait then re-read.""" import asyncio +from typing import Final +import fakeredis import pytest +from litellm.caching.redis_cache import RedisCache from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, ) +from litellm.proxy.mcp_server.outbound_credentials.redis_distributed_lock import RedisDistributedLock from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator import ( LockAcquisition, RedisRefreshCoordinator, ) from litellm.redis.stores.mcp_credential_refresh import RefreshLockTarget +from tests.unit.proxy.mcp_server.outbound_credentials.conftest import RedisClock, lost_write_reply _TARGET = RefreshLockTarget(user_id="u", server_id="s") @@ -67,7 +72,6 @@ async def test_winner_acquires_and_releases_with_the_same_token_and_never_reread async def reread(): reread_calls.append(1) - return None coord = RedisRefreshCoordinator(lock, new_token=lambda: "tok") result = await coord.run("u", "s", refresh, reread) @@ -140,7 +144,6 @@ async def test_loser_waits_for_the_holder_then_rereads_persisted_token(): async def refresh(): refresh_calls.append(1) - return None async def reread(): return OAuthToken(access_token="persisted-by-winner") @@ -219,7 +222,6 @@ async def test_loser_with_default_timeout_waits_past_the_lease_ttl_for_the_holde async def refresh(): refresh_calls.append(1) - return None async def reread(): return OAuthToken(access_token="persisted-by-winner") @@ -269,3 +271,251 @@ async def test_winner_stops_renewing_after_the_refresh_budget(): assert result is not None and result.access_token == "new" assert extends_during_budget == 4 # renewed at clock 5, 10, 15, 20 then stopped assert len(lock.extend_calls) == 4 # no further renewals after the budget + + +_KEY: Final = "ns:mcp:refresh_lock:u:s" + + +async def _never_refresh() -> OAuthToken | None: + pytest.fail("a loser must not refresh") + + +async def _never_reread() -> OAuthToken | None: + pytest.fail("a winner must not reread") + + +@pytest.mark.asyncio +async def test_winner_acquires_and_releases_with_the_same_token_and_never_rereads_with_redis_store( + refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis +) -> None: + refreshed: Final = OAuthToken(access_token="new") + + async def refresh() -> OAuthToken: + assert await refresh_raw.get(_KEY) == b"tok" + assert await refresh_raw.pttl(_KEY) == 10235 + return refreshed + + coord: Final = RedisRefreshCoordinator(refresh_lock, new_token=lambda: "tok", lock_ttl_seconds=10.2345) + assert await coord.run("u", "s", refresh, _never_reread) is refreshed + assert await refresh_raw.get(_KEY) is None + + +@pytest.mark.asyncio +async def test_winner_releases_even_when_refresh_raises_with_redis_store( + refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis +) -> None: + async def refresh() -> OAuthToken | None: + assert await refresh_raw.get(_KEY) == b"tok" + raise RuntimeError("boom") + + coord: Final = RedisRefreshCoordinator(refresh_lock, new_token=lambda: "tok") + with pytest.raises(RuntimeError, match="boom"): + await coord.run("u", "s", refresh, _never_reread) + assert await refresh_raw.get(_KEY) is None + + +@pytest.mark.asyncio +async def test_winner_renews_the_lock_while_refresh_runs_with_redis_store( + refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis, redis_clock: RedisClock +) -> None: + refreshed: Final = OAuthToken(access_token="new") + start: Final = redis_clock() + renewed: Final = asyncio.Event() + finish: Final = asyncio.Event() + + async def sleep(seconds: float) -> None: + assert seconds == 5.0 + redis_clock.advance(seconds) + if redis_clock() >= start + 10: + renewed.set() + await asyncio.Event().wait() + + async def refresh() -> OAuthToken: + await finish.wait() + return refreshed + + coord: Final = RedisRefreshCoordinator(refresh_lock, new_token=lambda: "tok", sleep=sleep, clock=redis_clock) + task: Final = asyncio.create_task(coord.run("u", "s", refresh, _never_reread)) + try: + await renewed.wait() + assert await refresh_raw.get(_KEY) == b"tok" + assert await refresh_raw.pttl(_KEY) == 5000 + finally: + finish.set() + assert await task is refreshed + assert await refresh_raw.get(_KEY) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("holder_seconds", (1.0, 15.0)) +async def test_loser_waits_for_the_holder_then_rereads_persisted_token_with_redis_store( + refresh_lock: RedisDistributedLock, + refresh_raw: fakeredis.FakeAsyncRedis, + redis_clock: RedisClock, + holder_seconds: float, +) -> None: + await refresh_raw.set(_KEY, "other-worker", px=30_000) + start: Final = redis_clock() + persisted: Final = OAuthToken(access_token="persisted-by-winner") + + async def sleep(seconds: float) -> None: + assert seconds == 0.05 + redis_clock.advance(0.5) + if redis_clock() >= start + holder_seconds: + await refresh_raw.delete(_KEY) + + async def reread() -> OAuthToken: + assert redis_clock() == start + holder_seconds + assert await refresh_raw.get(_KEY) is None + return persisted + + coord: Final = RedisRefreshCoordinator(refresh_lock, sleep=sleep, clock=redis_clock) + assert await coord.run("u", "s", _never_refresh, reread) is persisted + + +@pytest.mark.asyncio +async def test_loser_rereads_after_timeout_without_releasing_the_holder( + refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis, redis_clock: RedisClock +) -> None: + await refresh_raw.set(_KEY, "other-worker", px=10_000) + start: Final = redis_clock() + persisted: Final = OAuthToken(access_token="persisted") + + async def sleep(seconds: float) -> None: + redis_clock.advance(0.5) + + async def reread() -> OAuthToken: + assert redis_clock() == start + 1.0 + return persisted + + coord: Final = RedisRefreshCoordinator(refresh_lock, sleep=sleep, clock=redis_clock, wait_timeout_seconds=1.0) + assert await coord.run("u", "s", _never_refresh, reread) is persisted + assert await refresh_raw.get(_KEY) == b"other-worker" + assert await refresh_raw.pttl(_KEY) == 9000 + + +@pytest.mark.asyncio +async def test_lock_backend_error_refreshes_anyway_instead_of_serving_stale_with_redis_store( + refresh_server: fakeredis.FakeServer, refresh_lock: RedisDistributedLock +) -> None: + refresh_server.connected = False + refreshed: Final = OAuthToken(access_token="refreshed-despite-redis-down") + + async def refresh() -> OAuthToken: + return refreshed + + assert await RedisRefreshCoordinator(refresh_lock).run("u", "s", refresh, _never_reread) is refreshed + + +@pytest.mark.asyncio +async def test_unknown_applied_acquire_refreshes_without_releasing_an_unconfirmed_lease( + refresh_cache: RedisCache, refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis +) -> None: + client: Final = refresh_cache.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + client.set_response_callback("SET", lost_write_reply) + refreshed: Final = OAuthToken(access_token="new") + + async def refresh() -> OAuthToken: + assert await refresh_raw.get(_KEY) == b"tok" + return refreshed + + coord: Final = RedisRefreshCoordinator(refresh_lock, new_token=lambda: "tok") + assert await coord.run("u", "s", refresh, _never_reread) is refreshed + assert await refresh_raw.get(_KEY) == b"tok" + assert await refresh_raw.pttl(_KEY) == 10000 + + +@pytest.mark.asyncio +async def test_expired_owner_does_not_release_a_replacement_lease( + refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis, redis_clock: RedisClock +) -> None: + refreshed: Final = OAuthToken(access_token="new") + + async def refresh() -> OAuthToken: + assert await refresh_raw.get(_KEY) == b"tok" + redis_clock.advance(11.0) + assert await refresh_raw.set(_KEY, "replacement", nx=True, px=30_000) is True + return refreshed + + coord: Final = RedisRefreshCoordinator(refresh_lock, new_token=lambda: "tok", clock=redis_clock) + assert await coord.run("u", "s", refresh, _never_reread) is refreshed + assert await refresh_raw.get(_KEY) == b"replacement" + assert await refresh_raw.pttl(_KEY) == 30000 + + +@pytest.mark.asyncio +async def test_winner_stops_renewing_after_the_refresh_budget_with_redis_store( + refresh_lock: RedisDistributedLock, refresh_raw: fakeredis.FakeAsyncRedis, redis_clock: RedisClock +) -> None: + start: Final = redis_clock() + refreshed: Final = OAuthToken(access_token="new") + budget_reached: Final = asyncio.Event() + + async def sleep(seconds: float) -> None: + assert seconds == 5.0 + redis_clock.advance(seconds) + assert redis_clock() <= start + 20 + if redis_clock() == start + 20: + budget_reached.set() + + async def refresh() -> OAuthToken: + await budget_reached.wait() + for _ in range(100): + if await refresh_raw.pttl(_KEY) == 10000: + break + await asyncio.sleep(0) + else: + pytest.fail("the final renewal did not retain its lease tail") + redis_clock.advance(11.0) + assert await refresh_raw.get(_KEY) is None + return refreshed + + coord: Final = RedisRefreshCoordinator(refresh_lock, new_token=lambda: "tok", clock=redis_clock, sleep=sleep) + assert coord.wait_timeout_seconds > coord.refresh_budget_seconds + coord.lock_ttl_seconds + assert await coord.run("u", "s", refresh, _never_reread) is refreshed + assert await refresh_raw.get(_KEY) is None + + +@pytest.mark.asyncio +async def test_loser_rereads_if_redis_fails_while_waiting( + refresh_server: fakeredis.FakeServer, + refresh_lock: RedisDistributedLock, + refresh_raw: fakeredis.FakeAsyncRedis, + redis_clock: RedisClock, +) -> None: + await refresh_raw.set(_KEY, "other-worker", px=10_000) + persisted: Final = OAuthToken(access_token="persisted") + + async def sleep(seconds: float) -> None: + refresh_server.connected = False + + async def reread() -> OAuthToken: + return persisted + + coord: Final = RedisRefreshCoordinator(refresh_lock, sleep=sleep, clock=redis_clock) + assert await coord.run("u", "s", _never_refresh, reread) is persisted + refresh_server.connected = True + assert await refresh_raw.get(_KEY) == b"other-worker" + + +@pytest.mark.asyncio +async def test_runtime_coordinator_uses_the_auth_cache_redis_connection( + refresh_cache: RedisCache, refresh_raw: fakeredis.FakeAsyncRedis, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.caching.dual_cache import DualCache + from litellm.proxy import proxy_server + from litellm.proxy.mcp_server.outbound_credentials.runtime_refresh_coordinator import runtime_refresh_coordinator + + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=refresh_cache)) + coord: Final = runtime_refresh_coordinator() + assert coord is not None + refreshed: Final = OAuthToken(access_token="new") + + async def refresh() -> OAuthToken: + assert await refresh_raw.get(_KEY) is not None + assert await refresh_raw.pttl(_KEY) == 10000 + return refreshed + + assert await coord.run("u", "s", refresh, _never_reread) is refreshed + assert await refresh_raw.get(_KEY) is None