test(redis): retain refresh lock compatibility coverage

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()))

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

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