mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
refactor(redis-store): migrate MCP single-use markers
Co-authored-by: Yujong Lee <yujong@berri.ai>
This commit is contained in:
Devin AI
parent
01038214d3
commit
a42177a644
5 changed files with 208 additions and 216 deletions
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
41
litellm/redis/stores/mcp_single_use.py
Normal file
41
litellm/redis/stores/mcp_single_use.py
Normal file
|
|
@ -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}")
|
||||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue