From f9fa29b0b039937825fc30ca081f4c9ae769189e 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:38:31 +0000 Subject: [PATCH] test(redis): retain MCP marker failure and wire compatibility --- .../proxy/mcp_server/test_gateway_dcr_flow.py | 139 ++++++++++++++++++ 1 file changed, 139 insertions(+) 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 0e1dd5defb8..8228d9f2c5f 100644 --- a/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py @@ -1,8 +1,10 @@ """Tests for the aggregate gateway DCR flow (register, authorize, complete, token).""" +import asyncio import hashlib import json import re +import time from base64 import urlsafe_b64encode from datetime import datetime, timedelta, timezone from http.cookies import SimpleCookie @@ -72,6 +74,7 @@ from litellm.redis.stores.mcp_single_use import ( ConnectFlowUsed, RefreshFamilyRevoked, RefreshTokenUsed, + SingleUseMarker, ) from tests._support.fake_redis import fake_redis_cache @@ -825,6 +828,142 @@ async def test_single_use_guard_claims_each_marker_under_the_key_old_pods_use(ma LOOPBACK_REDIRECT_URI = "http://localhost:3118/callback" +_MARKER_CASES: Final = ( + (AuthCodeUsed("jti: é\n1"), "mcp_gateway_dcr_code_used:jti: é\n1"), + (ConnectFlowUsed("jti: é\n1"), "mcp_gateway_dcr_flow_used:jti: é\n1"), + (RefreshTokenUsed("jti: é\n1"), "mcp_gateway_dcr_refresh_used:jti: é\n1"), + (RefreshFamilyRevoked("jti: é\n1"), "mcp_gateway_dcr_refresh_family_revoked:jti: é\n1"), + ( + AssertionRedeemed("https://issuer.example/租户", "jti: é\n1"), + "mcp_gateway_dcr:id_jag:dabba8ce01737ba5cb47a28665f7bec742b072d198672c80bbc6e70a2781a221", + ), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("marker", "key"), _MARKER_CASES) +async def test_marker_families_interoperate_with_old_writers_and_preserve_expiry( + marker: SingleUseMarker, key: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard + + monkeypatch.setattr(time, "time", lambda: 1_700_000_000.0) + server: Final = fakeredis.FakeServer() + redis_cache: Final = fake_redis_cache(server, namespace="ns") + raw: Final = fakeredis.FakeAsyncRedis(server=server) + guard: Final = _SingleUseGuard(DualCache(redis_cache=redis_cache)) + wire_key: Final = f"ns:{key}" + + assert await guard.peek(marker) == "unclaimed" + assert await guard.claim(marker, 60) == "first" + assert await raw.keys("*") == [wire_key.encode()] + assert await raw.get(wire_key) == b"1" + assert await raw.pttl(wire_key) == 60_000 + assert await redis_cache.async_get_cache(key) == 1 + assert await redis_cache.async_increment(key, 1, ttl=5) == 2.0 + assert await guard.peek(marker) == "claimed" + assert await guard.claim(marker, 120) == "replayed" + assert await raw.get(wire_key) == b"3" + assert await raw.pttl(wire_key) == 60_000 + + assert await raw.delete(wire_key) == 1 + assert await redis_cache.async_increment(key, 1, ttl=60) == 1.0 + assert await raw.persist(wire_key) + assert await raw.pttl(wire_key) == -1 + other_worker: Final = _SingleUseGuard(DualCache(redis_cache=fake_redis_cache(server, namespace="ns"))) + assert await other_worker.peek(marker) == "claimed" + assert await other_worker.claim(marker, 17) == "replayed" + assert await raw.get(wire_key) == b"2" + assert await raw.pttl(wire_key) == 17_000 + + monkeypatch.setattr(time, "time", lambda: 1_700_000_018.0) + assert await guard.peek(marker) == "unclaimed" + claimed: Final = await asyncio.gather(guard.claim(marker, 30), other_worker.claim(marker, 30)) + assert sorted(claimed) == ["first", "replayed"] + assert await raw.get(wire_key) == b"2" + assert await raw.pttl(wire_key) == 30_000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("marker", "key"), _MARKER_CASES) +async def test_marker_families_use_canonical_keys_only_for_local_fallback(marker: SingleUseMarker, key: str) -> None: + from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard + + cache: Final = DualCache() + guard: Final = _SingleUseGuard(cache) + assert await guard.peek(marker) == "unclaimed" + assert await guard.claim(marker, 60) == "first" + assert cache.in_memory_cache.get_cache(key) == 1 + assert await guard.peek(marker) == "claimed" + assert await guard.claim(marker, 60) == "replayed" + assert cache.in_memory_cache.get_cache(key) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("marker", "key"), _MARKER_CASES) +async def test_marker_families_fail_closed_on_invalid_counter(marker: SingleUseMarker, key: str) -> None: + from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard + + server: Final = fakeredis.FakeServer() + cache: Final = DualCache(redis_cache=fake_redis_cache(server, namespace="ns")) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + guard: Final = _SingleUseGuard(cache) + assert await raw.set(f"ns:{key}", b"not a counter") + assert await guard.claim(marker, 60) == "unavailable" + assert await guard.peek(marker) == "claimed" + assert await raw.get(f"ns:{key}") == b"not a counter" + assert await raw.pttl(f"ns:{key}") == -1 + assert cache.in_memory_cache.get_cache(key) is None + + +class _RedisThatLosesClaimReply(fakeredis.FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> object: + await super().eval(script, numkeys, *keys_and_args) + raise RedisConnectionError("claim applied but reply was lost") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reply", (b"nan", b"invalid", 1)) +async def test_malformed_claim_reply_is_unavailable_without_local_fallback(reply: bytes | int) -> None: + from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard + + class RedisThatCorruptsClaimReply(fakeredis.FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> bytes | int: + await super().eval(script, numkeys, *keys_and_args) + return reply + + server: Final = fakeredis.FakeServer() + redis_cache: Final = fake_redis_cache(server, namespace="ns", client=RedisThatCorruptsClaimReply) + cache: Final = DualCache(redis_cache=redis_cache) + assert await _SingleUseGuard(cache).claim(AuthCodeUsed("invalid-reply"), 60) == "unavailable" + assert cache.in_memory_cache.get_cache("mcp_gateway_dcr_code_used:invalid-reply") is None + assert await fakeredis.FakeAsyncRedis(server=server).get("ns:mcp_gateway_dcr_code_used:invalid-reply") == b"1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("marker", "key"), _MARKER_CASES) +async def test_marker_families_refuse_replay_without_retrying_a_lost_reply( + marker: SingleUseMarker, key: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard + + monkeypatch.setattr(time, "time", lambda: 1_700_000_000.0) + server: Final = fakeredis.FakeServer() + redis_cache: Final = fake_redis_cache(server, namespace="ns", client=_RedisThatLosesClaimReply) + cache: Final = DualCache(redis_cache=redis_cache) + guard: Final = _SingleUseGuard(cache) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await guard.claim(marker, 60) == "unavailable" + assert await raw.get(f"ns:{key}") == b"1" + assert await raw.pttl(f"ns:{key}") == 60_000 + assert cache.in_memory_cache.get_cache(key) is None + other_worker: Final = _SingleUseGuard(DualCache(redis_cache=fake_redis_cache(server, namespace="ns"))) + assert await other_worker.claim(marker, 120) == "replayed" + assert await other_worker.peek(marker) == "claimed" + assert await raw.get(f"ns:{key}") == b"2" + assert await raw.pttl(f"ns:{key}") == 60_000 + + async def _complete(redirect_uri: str, delivery, cookies=None, handle=None, session_user_id="u1"): client_id = (await _register([redirect_uri]))["client_id"] if cookies is None: