From 6958731dd49637a32e56ccdb726d9cddfe7d6ffc Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 29 Oct 2026 18:35:05 -0700 Subject: [PATCH] refactor(redis-transport): read the MCP single-use marker through RedisTransport Co-Authored-By: Claude Opus 5.5 --- .../redis-access/references/transport.md | 3 + litellm/proxy/mcp_server/gateway_dcr_flow.py | 23 +++--- .../proxy/mcp_server/test_gateway_dcr_flow.py | 75 +++++++------------ 3 files changed, 43 insertions(+), 58 deletions(-) diff --git a/.agents/skills/redis-access/references/transport.md b/.agents/skills/redis-access/references/transport.md index 6d98a0c4135..a5f6394e39d 100644 --- a/.agents/skills/redis-access/references/transport.md +++ b/.agents/skills/redis-access/references/transport.md @@ -57,9 +57,12 @@ These rules are shared with `crates/redis-client` (its `AGENTS.md` has the Rust ## Moving raw access onto the transport +`peek` in `litellm/proxy/mcp_server/gateway_dcr_flow.py` is the reference. It went from `init_async_client().get(check_and_fix_namespace(key))` inside `try/except Exception` to a `match` on `redis_cache.transport().execute(Get(key))` + - Replace the raw client or script handle with transport calls, and delete hand-written namespacing and script registration - Keep the caller's outcome per failure. If it failed closed on any exception, it fails closed on any `RedisFailure` - Don't treat a `transport` failure on a write as "not applied" +- If the access is really one keyspace's logic (keys, codec, decoding), it belongs behind a store, see `store.md`. Only do that if asked, otherwise keep it a straight extraction ## Not on the transport yet diff --git a/litellm/proxy/mcp_server/gateway_dcr_flow.py b/litellm/proxy/mcp_server/gateway_dcr_flow.py index ae74d878b36..f2b1e66798a 100644 --- a/litellm/proxy/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/mcp_server/gateway_dcr_flow.py @@ -98,6 +98,7 @@ from litellm.proxy.mcp_server.outbound_credentials.session_token import ( open_session_refresh_token, open_session_token, ) +from litellm.redis.transport.contract import Get, RedisFailure from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -1393,21 +1394,21 @@ class _SingleUseGuard: async def peek(self, key: str) -> Literal["unclaimed", "claimed", "unavailable"]: """Read-only view of a single-use marker, resolved against the same shared authority as :meth:`claim` so introspection observes exactly the record redemption and revocation wrote. - A backend fault is ``"unavailable"`` (fail closed) rather than a guess either way, which is why - the read goes to the Redis client itself: the cache wrapper's ``async_get_cache`` turns a fault - into ``None``, and ``None`` here would pass as unclaimed. The key is namespaced exactly as the - wrapper namespaces every write, or a ``namespace`` configured under ``coordination_redis`` or - ``cache_params`` would hide the marker.""" + A backend fault is ``"unavailable"`` (fail closed) rather than a guess either way.""" from litellm.proxy.proxy_server import redis_usage_cache # noqa: PLC0415 # circular import at module load redis_cache: Final = redis_usage_cache or getattr(self._cache, "redis_cache", None) if redis_cache is not None: - try: - value = await redis_cache.init_async_client().get(redis_cache.check_and_fix_namespace(key=key)) - except Exception as e: # noqa: BLE001 # ANY Redis fault fails the read closed - verbose_logger.warning("mcp gateway single-use peek: shared cache backend unavailable: %s", e) - return "unavailable" - return "unclaimed" if value is None else "claimed" + match await redis_cache.transport().execute(Get(key)): + case RedisFailure() as failure: + verbose_logger.warning( + "mcp gateway single-use peek: shared cache backend unavailable: %s", failure.message + ) + return "unavailable" + case None: + return "unclaimed" + case _: + return "claimed" local: Final = await self._cache.async_get_cache(key, local_only=True) return "unclaimed" if local is None else "claimed" 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 a426cea131a..765d045ed09 100644 --- a/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py @@ -9,6 +9,7 @@ from http.cookies import SimpleCookie from typing import Final from urllib.parse import parse_qs, urlparse +import fakeredis import pytest from starlette.requests import Request from starlette.responses import Response @@ -63,6 +64,8 @@ from litellm.proxy.mcp_server.outbound_credentials.session_token import ( mint_session_refresh_token, mint_session_token, ) +from litellm.redis.transport.contract import Get, RedisFailure +from tests._support.fake_redis import fake_redis_cache MASTER_KEY = "sk-gateway-dcr-flow-tests" REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback" @@ -729,57 +732,35 @@ async def test_single_use_guard_fails_closed_when_redis_errors(): @pytest.mark.asyncio async def test_single_use_guard_peek_fails_closed_when_redis_errors(): - """The peek that gates a refresh on its chain's revocation marker must read the Redis client - itself: the cache wrapper's get swallows a fault into ``None``, which would make a revoked chain - look unclaimed for exactly as long as Redis is down.""" - from unittest.mock import AsyncMock, MagicMock - + """The peek that gates a refresh on its chain's revocation marker must tell a Redis fault from a + missing key, or a revoked chain looks unclaimed for exactly as long as Redis is down.""" from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard - cache = DualCache() - cache.redis_cache = MagicMock() - cache.redis_cache.async_get_cache = AsyncMock(return_value=None) - cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key) - cache.redis_cache.init_async_client.return_value.get = AsyncMock(side_effect=ConnectionError("redis down")) - + server = fakeredis.FakeServer() + cache = DualCache(redis_cache=fake_redis_cache(server)) guard = _SingleUseGuard(cache) + assert await guard.claim("family-fault", 60) == "first" + + server.connected = False assert await guard.peek("family-fault") == "unavailable" - cache.redis_cache.init_async_client.return_value.get = AsyncMock(return_value=b"1") + server.connected = True assert await guard.peek("family-fault") == "claimed" @pytest.mark.asyncio async def test_single_use_guard_peek_reads_the_key_under_the_namespace_claim_wrote(): - """A configured ``redis_namespace`` prefixes every key the cache wrapper writes, so the raw client - read behind peek must ask for the same prefixed key: a revocation written under ``ns:...`` and - peeked at the bare key would never be seen, and the revoked chain would keep renewing.""" - from unittest.mock import AsyncMock, MagicMock - + """A configured ``redis_namespace`` prefixes every key the cache writes, so peek must read the + same prefixed key: a revocation written under ``ns:...`` and peeked at the bare key would never + be seen, and the revoked chain would keep renewing.""" from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard - stored: dict[str, int] = {} # mutable-ok: the fake Redis store the test inspects + server = fakeredis.FakeServer() + guard = _SingleUseGuard(DualCache(redis_cache=fake_redis_cache(server, namespace="ns"))) - def namespaced(key: str) -> str: - return f"ns:{key}" - - async def increment(key: str, value: int, ttl: int) -> int: - stored[namespaced(key)] = stored.get(namespaced(key), 0) + value - return stored[namespaced(key)] - - async def raw_get(key: str) -> bytes | None: - return None if key not in stored else str(stored[key]).encode() - - cache = DualCache() - cache.redis_cache = MagicMock() - cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=namespaced) - cache.redis_cache.async_increment = AsyncMock(side_effect=increment) - cache.redis_cache.init_async_client.return_value.get = AsyncMock(side_effect=raw_get) - - guard = _SingleUseGuard(cache) assert await guard.peek("family-ns") == "unclaimed" assert await guard.claim("family-ns", 60) == "first" - assert set(stored) == {"ns:family-ns"} + assert await fakeredis.FakeAsyncRedis(server=server).keys("*") == [b"ns:family-ns"] assert await guard.peek("family-ns") == "claimed" @@ -2247,16 +2228,18 @@ async def test_revoke_burns_the_refresh_token_and_answers_200_for_dead_or_unknow assert garbage.status_code == 200 +_REDIS_DOWN: Final = RedisFailure(kind="transport", message="redis down") + + def _redis_that(async_increment, get=None): from unittest.mock import AsyncMock, MagicMock cache = DualCache() cache.redis_cache = MagicMock() cache.redis_cache.async_increment = async_increment - cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key) - cache.redis_cache.init_async_client.return_value.get = get or AsyncMock(return_value=None) + cache.redis_cache.transport.return_value.execute = get or AsyncMock(return_value=None) cache.redis_cache.async_get_cache = AsyncMock( - side_effect=AssertionError("peek must read the client, not the wrapper") + side_effect=AssertionError("peek must read through the transport, not the wrapper") ) cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory")) return cache @@ -2313,10 +2296,7 @@ class _RedisThatFaultsOnTheSecondWrite: self.landed: tuple[str, ...] = () self.writes = 0 - def check_and_fix_namespace(self, key: str) -> str: - return key - - def init_async_client(self) -> "_RedisThatFaultsOnTheSecondWrite": + def transport(self) -> "_RedisThatFaultsOnTheSecondWrite": return self async def async_increment(self, key: str, value: int, **kwargs: object) -> int: @@ -2326,8 +2306,9 @@ class _RedisThatFaultsOnTheSecondWrite: self.landed = (*self.landed, key) return self.landed.count(key) - async def get(self, key: str) -> int | None: - return self.landed.count(key) or None + async def execute(self, command: Get) -> bytes | None: + count: Final = self.landed.count(command.key) + return str(count).encode() if count else None @pytest.mark.asyncio @@ -2421,7 +2402,7 @@ async def test_refresh_of_a_rotated_token_answers_503_before_minting_while_redis minter = _Minter() claim = AsyncMock(return_value=1) unreadable = await _refresh_native( - rotated, client_id, minter, _redis_that(claim, get=AsyncMock(side_effect=ConnectionError("redis down"))) + rotated, client_id, minter, _redis_that(claim, get=AsyncMock(return_value=_REDIS_DOWN)) ) assert unreadable.status_code == 503 assert json.loads(unreadable.body)["error"] == "temporarily_unavailable" @@ -2728,7 +2709,7 @@ async def test_introspect_503s_while_the_single_use_record_cannot_be_read(): keys, now, principal = _introspection_fixtures() minted = mint_session_refresh_token(principal, keys, now) - unreadable = _redis_that(AsyncMock(return_value=1), get=AsyncMock(side_effect=ConnectionError("redis down"))) + unreadable = _redis_that(AsyncMock(return_value=1), get=AsyncMock(return_value=_REDIS_DOWN)) status, body = await _introspect(minted.token.get_secret_value(), cache=unreadable) assert (status, body["error"]) == (503, "temporarily_unavailable")