diff --git a/litellm/proxy/mcp_server/gateway_dcr_flow.py b/litellm/proxy/mcp_server/gateway_dcr_flow.py index ae74d878b36..debf2194d98 100644 --- a/litellm/proxy/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/mcp_server/gateway_dcr_flow.py @@ -44,7 +44,7 @@ import secrets from base64 import urlsafe_b64encode from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import Final, Literal, Protocol, TypeVar from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse @@ -98,6 +98,8 @@ from litellm.proxy.mcp_server.outbound_credentials.session_token import ( open_session_refresh_token, open_session_token, ) +from litellm.redis.stores.mcp_single_use import TransportMcpSingleUseStore +from litellm.redis.transport.contract import RedisFailure from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -1348,7 +1350,7 @@ class _SingleUseGuard: id replays through a different worker): - When a Redis backend is configured it is the SOLE authority: the claim goes straight to Redis - (``INCR`` is atomic across replicas), and any Redis fault fails the claim CLOSED — it never falls + (``INCRBYFLOAT`` is atomic across replicas), and any Redis fault fails the claim CLOSED — it never falls back to the per-worker in-memory count (``DualCache.async_increment_cache`` does fall back, which is exactly the replay window this avoids). - With no Redis configured (single-replica) the in-memory increment is authoritative within the one @@ -1374,40 +1376,35 @@ class _SingleUseGuard: # proxy already treats as cross-worker, so no call site can wire the guarantee away. redis_cache: Final = redis_usage_cache or getattr(self._cache, "redis_cache", None) if redis_cache is not None: - # Shared, atomic authority for multi-replica deployments. Claim ONLY against Redis and fail - # CLOSED on any Redis fault (async_increment re-raises) rather than fall back to the - # per-worker in-memory count, which would let each replica observe count==1 and replay the id. - try: - count = await redis_cache.async_increment(key, 1, ttl=ttl_seconds) - except Exception as e: # noqa: BLE001 # ANY Redis fault fails the single-use claim closed - verbose_logger.warning( - "mcp gateway single-use claim: shared cache backend unavailable, failing closed: %s", e - ) - return "unavailable" - return "first" if count == 1 else "replayed" + match await TransportMcpSingleUseStore(redis_cache.transport()).claim( + key, ttl=timedelta(seconds=ttl_seconds) + ): + case RedisFailure(): + return "unavailable" + case 1: + return "first" + case _: + return "replayed" # No shared backend configured (single-replica): the in-memory increment is authoritative. - count = await self._cache.async_increment_cache(key, 1, ttl=ttl_seconds, local_only=True) + count: Final = await self._cache.async_increment_cache(key, 1, ttl=ttl_seconds, local_only=True) return "first" if count == 1 else "replayed" @with_service_target(_DCR_CLAIMS_TARGET) 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 TransportMcpSingleUseStore(redis_cache.transport()).peek(key): + case RedisFailure(): + return "unavailable" + case True: + return "claimed" + case False: + return "unclaimed" local: Final = await self._cache.async_get_cache(key, local_only=True) return "unclaimed" if local is None else "claimed" diff --git a/litellm/redis/stores/mcp_single_use.py b/litellm/redis/stores/mcp_single_use.py new file mode 100644 index 00000000000..31d754732eb --- /dev/null +++ b/litellm/redis/stores/mcp_single_use.py @@ -0,0 +1,41 @@ +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from litellm.redis.transport.contract import EvalReply, RedisFailure, RedisTransport + +_CLAIM: Final = """ +local count = redis.call("incrbyfloat", KEYS[1], ARGV[1]) +if redis.call("ttl", KEYS[1]) < 0 then + redis.call("expire", KEYS[1], ARGV[2]) +end +return count +""" + + +class McpSingleUseStore(Protocol): + async def claim(self, key: str, *, ttl: timedelta) -> int | RedisFailure: ... + + async def peek(self, key: str) -> bool | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportMcpSingleUseStore: + transport: RedisTransport + + async def claim(self, key: str, *, ttl: timedelta) -> int | RedisFailure: + reply: Final = await self.transport.eval(_CLAIM, (key,), (1, ttl // timedelta(seconds=1))) + return reply if isinstance(reply, RedisFailure) else _count(reply) + + async def peek(self, key: str) -> bool | RedisFailure: + reply: Final = await self.transport.get(key) + return reply if isinstance(reply, RedisFailure) else reply is not None + + +def _count(reply: EvalReply) -> int | RedisFailure: + if not isinstance(reply, bytes): + return RedisFailure(kind="decode", message=f"single-use count is not a bulk string: {reply!r}") + try: + return int(float(reply)) + except ValueError: + return RedisFailure(kind="decode", message=f"single-use count is not a number: {reply!r}") diff --git a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py index 8d4571c2f01..55b64af70cc 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -14,12 +14,14 @@ from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import fakeredis import httpx import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import ProxyException from litellm.proxy.anthropic_endpoints import gateway_endpoints from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken @@ -72,55 +74,18 @@ _COMPLETED_SESSION: Final = MappingProxyType( ) -class _SharedRedisFake: - """A Redis double with a configured namespace, prefixing every key its wrapper methods touch the - way ``RedisCache`` does, while its raw client answers only the exact key it is asked for.""" - - namespace: Final = "qa" - +class _SharedRedisCache(RedisCache): def __init__(self) -> None: - self.values: Mapping[str, object] = MappingProxyType({}) - self.counters: Mapping[str, float] = MappingProxyType({}) + self._server: Final = fakeredis.FakeServer() + super().__init__(host="127.0.0.1", port=1, namespace="qa", socket_timeout=0.1) + self.redis_client = fakeredis.FakeRedis(server=self._server) - def check_and_fix_namespace(self, key: str) -> str: - return key if key.startswith(f"{self.namespace}:") else f"{self.namespace}:{key}" - - def set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.values = MappingProxyType({**self.values, self.check_and_fix_namespace(key): value}) - - def get_cache(self, key: str, **kwargs: object) -> object: - return self.values.get(self.check_and_fix_namespace(key)) - - def delete_cache(self, key: str) -> None: - gone: Final = self.check_and_fix_namespace(key) - self.values = MappingProxyType({name: value for name, value in self.values.items() if name != gone}) - - async def async_delete_cache(self, key: str) -> None: - self.delete_cache(key) - - async def async_increment(self, key: str, value: float, **kwargs: object) -> float: - namespaced: Final = self.check_and_fix_namespace(key) - incremented: Final = self.counters.get(namespaced, 0) + value - self.counters = MappingProxyType({**self.counters, namespaced: incremented}) - return incremented - - async def async_get_cache(self, key: str, **kwargs: object) -> object: - return self.values.get(self.check_and_fix_namespace(key)) - - def init_async_client(self) -> "_SharedRedisClientFake": - return _SharedRedisClientFake(self) + def init_async_client(self) -> fakeredis.FakeAsyncRedis: + return fakeredis.FakeAsyncRedis(server=self._server) -class _SharedRedisClientFake: - def __init__(self, store: _SharedRedisFake) -> None: - self._store = store - - async def get(self, key: str) -> object: - return self._store.counters.get(key, self._store.values.get(key)) - - -def _replica(redis: _SharedRedisFake) -> DualCache: - return DualCache(redis_cache=redis, default_in_memory_ttl=600) # pyright: ignore[reportArgumentType] # duck-typed Redis double +def _replica(redis: RedisCache) -> DualCache: + return DualCache(redis_cache=redis, default_in_memory_ttl=600) class _Minter: @@ -525,7 +490,7 @@ def test_token_unknown_team_grants_is_invalid_grant(): def test_token_mints_on_a_replica_that_did_not_start_the_login(): - redis: Final = _SharedRedisFake() + redis: Final = _SharedRedisCache() _set_cli_sso_flow(login_id=_SHARED_LOGIN_ID, cache=_replica(redis), flow=_completed_flow()) with _gateway_env(cache=_replica(redis)) as (client, _), patch(_MINT, return_value="sk-session") as mint: @@ -537,7 +502,7 @@ def test_token_mints_on_a_replica_that_did_not_start_the_login(): def test_token_refuses_a_device_code_another_replica_already_claimed(): - redis: Final = _SharedRedisFake() + redis: Final = _SharedRedisCache() replica_a: Final = _replica(redis) _set_cli_sso_flow(login_id=_SHARED_LOGIN_ID, cache=replica_a, flow=_completed_flow()) assert asyncio.run(gateway_endpoints._claim_device_code(_SHARED_LOGIN_ID, replica_a)) is True @@ -616,7 +581,7 @@ def test_refresh_grant_mints_a_new_bearer_and_rotates_the_refresh_token(): def test_refresh_grant_replay_is_refused_on_a_replica_that_did_not_serve_the_rotation(): - redis: Final = _SharedRedisFake() + redis: Final = _SharedRedisCache() minter: Final = _Minter() with _gateway_env(claim_cache=_replica(redis), minter=minter) as (client, cache): refresh_token = str(_signed_in(client, cache)["refresh_token"]) @@ -761,7 +726,7 @@ def test_revoke_on_one_replica_ends_the_chain_on_another_through_a_namespaced_re with the configured ``redis_namespace``; the replica serving the next renewal must look for the marker under that same prefix, or the revoked chain keeps renewing everywhere but where it signed out.""" - redis: Final = _SharedRedisFake() + redis: Final = _SharedRedisCache() minter: Final = _Minter() with _gateway_env(claim_cache=_replica(redis), minter=minter) as (client, cache): stored = str(_signed_in(client, cache)["refresh_token"]) 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..97e3e6b8daa 100644 --- a/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py @@ -9,10 +9,13 @@ from http.cookies import SimpleCookie from typing import Final from urllib.parse import parse_qs, urlparse +import fakeredis import pytest +from redis.exceptions import ConnectionError as RedisConnectionError from starlette.requests import Request from starlette.responses import Response +import litellm from litellm.caching.caching import DualCache from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.mcp_server.gateway_dcr_flow import ( @@ -63,6 +66,7 @@ from litellm.proxy.mcp_server.outbound_credentials.session_token import ( mint_session_refresh_token, mint_session_token, ) +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" @@ -691,95 +695,98 @@ async def test_single_use_guard_in_memory_is_single_use_within_process(): @pytest.mark.asyncio async def test_single_use_guard_uses_redis_as_sole_authority_when_configured(): - """With Redis configured it is the SOLE authority: the shared INCR result decides the claim (1 → - first caller, >1 → replay), and the per-worker in-memory count is never consulted.""" - from unittest.mock import AsyncMock, MagicMock - + """With Redis configured it is the SOLE authority: the shared count decides the claim (1 is the + first caller, more is a replay), and the per-worker in-memory count is never written.""" from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard - cache = DualCache() - cache.redis_cache = MagicMock() - cache.redis_cache.async_increment = AsyncMock(return_value=1) - # in-memory must NOT be consulted when Redis is configured — poison it so any fallback is visible. - cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory")) - + server = fakeredis.FakeServer() + cache = DualCache(redis_cache=fake_redis_cache(server)) guard = _SingleUseGuard(cache) + assert await guard.claim("jti-redis", 60) == "first" - cache.redis_cache.async_increment = AsyncMock(return_value=2) assert await guard.claim("jti-redis", 60) == "replayed" + assert await _SingleUseGuard(DualCache(redis_cache=fake_redis_cache(server))).claim("jti-redis", 60) == "replayed" + assert cache.in_memory_cache.get_cache("jti-redis") is None + + +@pytest.mark.asyncio +async def test_single_use_guard_writes_the_marker_the_old_increment_wrote_with_its_ttl(): + """Old and new pods share the marker during a rolling deploy, so a claim writes the same + INCRBYFLOAT string with the same TTL, and each side counts the other's claim.""" + from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard + + server = fakeredis.FakeServer() + redis_cache = fake_redis_cache(server, namespace="ns") + raw = fakeredis.FakeAsyncRedis(server=server) + guard = _SingleUseGuard(DualCache(redis_cache=redis_cache)) + + assert await guard.claim("new-first", 60) == "first" + assert await raw.get("ns:new-first") == b"1" + assert 59_000 < await raw.pttl("ns:new-first") <= 60_000 + assert await redis_cache.async_increment("new-first", 1, ttl=60) == 2.0 + + assert await redis_cache.async_increment("old-first", 1, ttl=60) == 1.0 + assert await guard.claim("old-first", 60) == "replayed" + assert await guard.peek("old-first") == "claimed" + + +@pytest.mark.asyncio +async def test_single_use_guard_keeps_the_ttl_a_marker_already_has(): + from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard + + server = fakeredis.FakeServer() + raw = fakeredis.FakeAsyncRedis(server=server) + guard = _SingleUseGuard(DualCache(redis_cache=fake_redis_cache(server))) + + assert await guard.claim("jti-ttl", 600) == "first" + assert await guard.claim("jti-ttl", 5) == "replayed" + assert 590_000 < await raw.pttl("jti-ttl") <= 600_000 @pytest.mark.asyncio async def test_single_use_guard_fails_closed_when_redis_errors(): """A Redis fault must fail the claim CLOSED (refuse the id) rather than fall back to the per-worker - in-memory count — which would let each replica observe count==1 and replay the one-time id (the - Cursor/Veria replay-across-workers finding).""" - from unittest.mock import AsyncMock, MagicMock - + in-memory count, which would let each replica observe count==1 and replay the one-time id.""" from litellm.proxy.mcp_server.gateway_dcr_flow import _SingleUseGuard - cache = DualCache() - cache.redis_cache = MagicMock() - cache.redis_cache.async_increment = AsyncMock(side_effect=ConnectionError("redis down")) - cache.async_increment_cache = AsyncMock(return_value=1) # would fail OPEN if the guard fell back + server = fakeredis.FakeServer() + cache = DualCache(redis_cache=fake_redis_cache(server)) + server.connected = False - guard = _SingleUseGuard(cache) - assert await guard.claim("jti-fault", 60) == "unavailable" # fail closed, not a fallback count of 1 + assert await _SingleUseGuard(cache).claim("jti-fault", 60) == "unavailable" + assert cache.in_memory_cache.get_cache("jti-fault") is None @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() + guard = _SingleUseGuard(DualCache(redis_cache=fake_redis_cache(server))) + assert await guard.claim("family-fault", 60) == "first" - guard = _SingleUseGuard(cache) + 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,38 +2254,25 @@ async def test_revoke_burns_the_refresh_token_and_answers_200_for_dead_or_unknow assert garbage.status_code == 200 -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.async_get_cache = AsyncMock( - side_effect=AssertionError("peek must read the client, not the wrapper") - ) - cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory")) - return cache +def _shared_redis(server: fakeredis.FakeServer) -> DualCache: + return DualCache(redis_cache=fake_redis_cache(server)) @pytest.mark.asyncio async def test_revoke_answers_503_while_the_shared_record_cannot_be_written_then_burns_the_token(): """A revocation whose single-use marker never reached Redis must not report success: the token is still redeemable on every worker, so the client has to hear 503 (RFC 7009 2.2.1) and retry.""" - from unittest.mock import AsyncMock - client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] issued = DualCache() payload = json.loads( (await _redeem_native(await _native_code(client_id, cache=issued), client_id, _Minter(), cache=issued)).body ) + server = fakeredis.FakeServer() + shared = _shared_redis(server) + server.connected = False redis_down = await revoke_refresh_token( - token=payload["refresh_token"], - client_id=client_id, - master_key=MASTER_KEY, - cache=_redis_that(AsyncMock(side_effect=ConnectionError("redis down"))), + token=payload["refresh_token"], client_id=client_id, master_key=MASTER_KEY, cache=shared ) assert redis_down.status_code == 503 assert json.loads(redis_down.body) == { @@ -2287,47 +2281,32 @@ async def test_revoke_answers_503_while_the_shared_record_cannot_be_written_then } assert redis_down.headers["cache-control"] == "no-store" + server.connected = True redis_back = await revoke_refresh_token( - token=payload["refresh_token"], - client_id=client_id, - master_key=MASTER_KEY, - cache=_redis_that(AsyncMock(return_value=1)), + token=payload["refresh_token"], client_id=client_id, master_key=MASTER_KEY, cache=shared ) assert redis_back.status_code == 200 assert json.loads(redis_back.body) == {} already_burned = await revoke_refresh_token( - token=payload["refresh_token"], - client_id=client_id, - master_key=MASTER_KEY, - cache=_redis_that(AsyncMock(return_value=2)), + token=payload["refresh_token"], client_id=client_id, master_key=MASTER_KEY, cache=shared ) assert already_burned.status_code == 200 -class _RedisThatFaultsOnTheSecondWrite: - """A shared Redis double whose first increment lands and whose second raises, the fault between - revoke's two writes; its raw client reads back exactly the keys that landed, and later writes land.""" +class _RedisThatFaultsOnTheSecondClaim(fakeredis.FakeAsyncRedis): + """A shared Redis whose second single-use claim faults, the fault between revoke's two writes; + every other command, and every later claim, lands.""" - def __init__(self) -> None: - self.landed: tuple[str, ...] = () - self.writes = 0 + def __init__(self, server: fakeredis.FakeServer) -> None: + super().__init__(server=server) + self.claims = 0 - def check_and_fix_namespace(self, key: str) -> str: - return key - - def init_async_client(self) -> "_RedisThatFaultsOnTheSecondWrite": - return self - - async def async_increment(self, key: str, value: int, **kwargs: object) -> int: - self.writes += 1 - if self.writes == 2: - raise ConnectionError("redis down") - 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 evalsha(self, sha: 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) @pytest.mark.asyncio @@ -2343,8 +2322,13 @@ async def test_revoke_ends_the_chain_first_so_a_fault_before_the_jti_burn_leaves copied_and_rotated = json.loads( (await _refresh_native(payload["refresh_token"], client_id, _Minter(), issued)).body ) - redis: Final = _RedisThatFaultsOnTheSecondWrite() - shared: Final = DualCache(redis_cache=redis) # pyright: ignore[reportArgumentType] # duck-typed Redis double + server: Final = fakeredis.FakeServer() + redis_cache: Final = fake_redis_cache(server) + litellm.in_memory_llm_clients_cache.set_cache( + key=redis_cache._get_async_client_cache_key(), # pyright: ignore[reportPrivateUsage] # seeds the per-loop client slot RedisCache reads + value=_RedisThatFaultsOnTheSecondClaim(server), + ) + shared: Final = DualCache(redis_cache=redis_cache) half_written = await revoke_refresh_token( token=payload["refresh_token"], client_id=client_id, master_key=MASTER_KEY, cache=shared @@ -2368,33 +2352,26 @@ async def test_refresh_answers_503_without_burning_the_token_while_redis_is_down """Fail closed, but say why: a refresh the shared backend could not record is refused with 503 rather than ``invalid_grant``, so the CLI keeps the key it has and retries instead of telling the user the token was already used and sending them back through the browser.""" - from unittest.mock import AsyncMock - client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] issued = DualCache() payload = json.loads( (await _redeem_native(await _native_code(client_id, cache=issued), client_id, _Minter(), cache=issued)).body ) + server = fakeredis.FakeServer() + shared = _shared_redis(server) - redis_down = await _refresh_native( - payload["refresh_token"], - client_id, - _Minter(), - _redis_that(AsyncMock(side_effect=ConnectionError("redis down"))), - ) + server.connected = False + redis_down = await _refresh_native(payload["refresh_token"], client_id, _Minter(), shared) assert redis_down.status_code == 503 assert json.loads(redis_down.body)["error"] == "temporarily_unavailable" assert "refresh_token" not in json.loads(redis_down.body) - redis_back = await _refresh_native( - payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)) - ) + server.connected = True + redis_back = await _refresh_native(payload["refresh_token"], client_id, _Minter(), shared) assert redis_back.status_code == 200 assert json.loads(redis_back.body)["refresh_token"] != payload["refresh_token"] - replayed = await _refresh_native( - payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=2)) - ) + replayed = await _refresh_native(payload["refresh_token"], client_id, _Minter(), shared) assert replayed.status_code == 400 assert json.loads(replayed.body)["error"] == "invalid_grant" @@ -2403,32 +2380,26 @@ async def test_refresh_answers_503_without_burning_the_token_while_redis_is_down async def test_refresh_of_a_rotated_token_answers_503_before_minting_while_redis_cannot_be_read(): """A rotated token's chain marker is read before anything is minted or claimed; a Redis read fault is a 503, never a pass: the token stays unburned and unminted until Redis answers.""" - from unittest.mock import AsyncMock - client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] issued = DualCache() payload = json.loads( (await _redeem_native(await _native_code(client_id, cache=issued), client_id, _Minter(), cache=issued)).body ) - rotated = json.loads( - ( - await _refresh_native( - payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)) - ) - ).body - )["refresh_token"] + server = fakeredis.FakeServer() + shared = _shared_redis(server) + rotated = json.loads((await _refresh_native(payload["refresh_token"], client_id, _Minter(), shared)).body)[ + "refresh_token" + ] 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"))) - ) + server.connected = False + unreadable = await _refresh_native(rotated, client_id, minter, shared) assert unreadable.status_code == 503 assert json.loads(unreadable.body)["error"] == "temporarily_unavailable" assert minter.calls == [] - assert claim.await_count == 0 - readable = await _refresh_native(rotated, client_id, _Minter(), _redis_that(AsyncMock(return_value=1))) + server.connected = True + readable = await _refresh_native(rotated, client_id, _Minter(), shared) assert readable.status_code == 200 assert json.loads(readable.body)["refresh_token"] != rotated @@ -2724,15 +2695,17 @@ async def test_introspect_fails_closed_on_dead_user_and_503s_on_outage(): @pytest.mark.asyncio async def test_introspect_503s_while_the_single_use_record_cannot_be_read(): """A token whose chain may have been revoked is never reported active on a Redis read fault.""" - from unittest.mock import AsyncMock - 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"))) - status, body = await _introspect(minted.token.get_secret_value(), cache=unreadable) + server = fakeredis.FakeServer() + shared = _shared_redis(server) + + server.connected = False + status, body = await _introspect(minted.token.get_secret_value(), cache=shared) assert (status, body["error"]) == (503, "temporarily_unavailable") - status, body = await _introspect(minted.token.get_secret_value(), cache=_redis_that(AsyncMock(return_value=1))) + server.connected = True + status, body = await _introspect(minted.token.get_secret_value(), cache=shared) assert (status, body["active"]) == (200, True) diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 6003a21ff57..2155a75858d 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -201,7 +201,7 @@ def test_a_targeted_coroutine_closed_from_another_context_does_not_raise() -> No assert current_service_target() is None -_DIRECT_REDIS_CALL: Final = re.compile(r"\b_?redis_cache\.(?!async_register_script\b)(?:async_)?\w+\(") +_DIRECT_REDIS_CALL: Final = re.compile(r"\b_?redis_cache\.(?!(?:async_register_script|transport)\b)(?:async_)?\w+\(") @dataclass(frozen=True, slots=True) @@ -226,6 +226,22 @@ def _scan_function(source: str, fn: ast.FunctionDef | ast.AsyncFunctionDef) -> _ ) +@pytest.mark.parametrize( + ("body", "reaches_redis"), + ( + ("redis_cache.transport()", False), + ('redis_cache.async_register_script("lua")', False), + ('redis_cache.async_get_cache("key")', True), + ('redis_cache.transport(); redis_cache.async_get_cache("key")', True), + ), +) +def test_redis_family_scan_distinguishes_factories_from_commands(body: str, reaches_redis: bool) -> None: + source: Final = f"def call():\n {body}\n" + function: Final = ast.parse(source).body[0] + assert isinstance(function, ast.FunctionDef) + assert _scan_function(source, function).reaches_redis_directly is reaches_redis + + def _covered_by_callers(scans: tuple[_FunctionScan, ...], covered: frozenset[str]) -> frozenset[str]: """Close ``covered`` over functions whose every in-file caller already declares a family.""" callers: Final = {