scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 360

refactor(redis): move MCP per-user tokens into a typed store

Browse source
This commit is contained in:
Devin AI 2026-10-30 21:28:32 +00:00
parent dd35d6f0d3
commit 3e0dda2017
7 changed files with 482 additions and 175 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/proxy/mcp_server/oauth2_token_cache.py
View file

@ -17,13 +17,13 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
AUTH_OBJECTS_TARGET,
MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE,
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_PER_USER_TOKEN_DEFAULT_TTL,
MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
@ -34,9 +34,11 @@ from litellm.proxy.mcp_server.oauth_utils import (
build_upstream_oauth2_token_request,
resolve_upstream_resource,
)
from litellm.proxy.mcp_server.outbound_credentials.dual_cache_token_backend import DualCacheTokenCacheBackend
from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
from litellm.proxy.mcp_server.outbound_credentials.token_cache_codec import OAuthTokenCacheCodec
from litellm.proxy.mcp_server.utils import MCP_OAUTH_TOKENS_TARGET
from litellm.redis.stores.mcp_per_user_tokens import per_user_token_key
from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
@ -241,7 +243,7 @@ class MCPPerUserTokenCache:
"""
def _cache_key(self, user_id: str, server_id: str) -> str:
return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}"
return per_user_token_key(user_id, server_id)
def _codec(self) -> OAuthTokenCacheCodec:
return OAuthTokenCacheCodec(
@ -258,11 +260,7 @@ class MCPPerUserTokenCache:
try:
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
key: Final = self._cache_key(user_id, server_id)
encrypted: Final = await user_api_key_cache.async_get_cache(key)
if encrypted is None:
return None
return self._codec().decode(encrypted)
return await DualCacheTokenCacheBackend(user_api_key_cache, self._codec()).get(user_id, server_id)
except Exception as exc:
verbose_logger.debug(
"MCPPerUserTokenCache.get failed for user=%s server=%s: %s",
@ -285,11 +283,12 @@ class MCPPerUserTokenCache:
try:
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
key: Final = self._cache_key(user_id, server_id)
encrypted: Final = self._codec().encode(
OAuthToken(access_token=access_token, identity_binding_proof=identity_binding_proof)
await DualCacheTokenCacheBackend(user_api_key_cache, self._codec()).cache_token(
user_id,
server_id,
OAuthToken(access_token=access_token, identity_binding_proof=identity_binding_proof),
ttl,
)
await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl)
verbose_logger.debug(
"MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds",
user_id,
@ -304,16 +303,18 @@ class MCPPerUserTokenCache:
exc,
)
@with_service_target(AUTH_OBJECTS_TARGET)
async def delete(self, user_id: str, server_id: str) -> None:
"""Invalidate the cached token in Redis, here, and in every peer worker's in-memory layer."""
try:
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( # noqa: PLC0415 # proxy import cycle
evict_and_broadcast,
publish_auth_cache_invalidation,
)
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
key: Final = self._cache_key(user_id, server_id)
await evict_and_broadcast((key,), user_api_key_cache)
await DualCacheTokenCacheBackend(user_api_key_cache, self._codec()).evict_token(user_id, server_id)
await publish_auth_cache_invalidation(cache_key=key)
except Exception as exc:
verbose_logger.debug(
"MCPPerUserTokenCache.delete failed for user=%s server=%s: %s",

87
litellm/proxy/mcp_server/outbound_credentials/dual_cache_token_backend.py
View file

@ -10,10 +10,14 @@ the other across the cutover). A missing or undecryptable entry reads as a miss.
from __future__ import annotations
from dataclasses import KW_ONLY, dataclass
from datetime import timedelta
from typing import Final, Protocol
from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger
from litellm.caching.redis_batch import active_request_redis_batch
from litellm.caching.redis_cache import RedisCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import (
OAuthToken,
)
@ -21,16 +25,34 @@ from litellm.proxy.mcp_server.outbound_credentials.token_cache_codec import (
OAuthTokenCacheCodec,
)
from litellm.proxy.mcp_server.utils import MCP_OAUTH_TOKENS_TARGET
from litellm.redis.stores.mcp_per_user_tokens import (
TOKEN_KEY_PREFIX,
McpPerUserTokenStore,
TransportMcpPerUserTokenStore,
per_user_token_key,
)
from litellm.redis.transport.contract import RedisFailure
class TokenMemory(Protocol):
async def async_get_cache(self, key: str) -> object: ...
async def async_set_cache(self, key: str, value: str, *, ttl: float | None = None) -> None: ...
def delete_cache(self, key: str) -> None: ...
class AsyncCache(Protocol):
"""The slice of LiteLLM's ``DualCache`` this backend needs (Redis-backed, shared across workers)."""
async def async_get_cache(self, key: str) -> object | None: ...
@property
def redis_cache(self) -> RedisCache | None: ...
async def async_set_cache(self, key: str, value: str, ttl: float | None = None) -> None: ...
@property
def in_memory_cache(self) -> TokenMemory: ...
async def async_delete_cache(self, key: str) -> None: ...
@property
def default_in_memory_ttl(self) -> float | None: ...
@dataclass(frozen=True, slots=True)
@ -44,17 +66,39 @@ class DualCacheTokenCacheBackend:
cache: AsyncCache
codec: OAuthTokenCacheCodec
_: KW_ONLY
key_prefix: str = "mcp:per_user_token:"
key_prefix: str = TOKEN_KEY_PREFIX
def _key(self, user_id: str, server_id: str) -> str:
return f"{self.key_prefix}{user_id}:{server_id}"
return per_user_token_key(user_id, server_id, prefix=self.key_prefix)
def _store(self) -> McpPerUserTokenStore | None:
redis: Final = self.cache.redis_cache
return None if redis is None else TransportMcpPerUserTokenStore(redis.transport(), self.key_prefix)
async def _read(self, user_id: str, server_id: str) -> object:
key: Final = self._key(user_id, server_id)
local: Final = await self.cache.in_memory_cache.async_get_cache(key)
store: Final = self._store()
if local is not None or store is None:
return local
redis: Final = self.cache.redis_cache
batch: Final = None if redis is None else active_request_redis_batch(redis)
if batch is not None and batch.read_as_missing(key):
return None
blob: Final = await store.read(user_id, server_id)
if isinstance(blob, RedisFailure):
return None
if blob is not None:
ttl: Final = self.cache.default_in_memory_ttl
await self.cache.in_memory_cache.async_set_cache(key, blob, **({"ttl": ttl} if ttl is not None else {}))
return blob
@with_service_target(MCP_OAUTH_TOKENS_TARGET)
async def get(self, user_id: str, server_id: str) -> OAuthToken | None:
try:
blob: Final = await self.cache.async_get_cache(self._key(user_id, server_id))
blob: Final = await self._read(user_id, server_id)
return self.codec.decode(blob) if isinstance(blob, str) else None
except Exception as exc: # noqa: BLE001
except Exception as exc: # noqa: BLE001 # cache and codec failures remain safe misses
verbose_logger.debug("MCP per-user token cache get failed (miss): %s", exc)
return None
@ -62,18 +106,37 @@ class DualCacheTokenCacheBackend:
async def set(self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float) -> None:
if ttl_seconds <= 0:
return
await self.cache_token(user_id, server_id, token, ttl_seconds)
async def cache_token(self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float) -> None:
try:
await self.cache.async_set_cache(
ciphertext: Final = self.codec.encode(token)
await self.cache.in_memory_cache.async_set_cache(
self._key(user_id, server_id),
self.codec.encode(token),
ciphertext,
ttl=ttl_seconds,
)
except Exception as exc: # noqa: BLE001
store: Final = self._store()
if store is not None:
await store.write(
user_id,
server_id,
ciphertext,
ttl=timedelta(seconds=int(ttl_seconds)),
defer=isinstance(self.cache, UserApiKeyCache),
)
except Exception as exc: # noqa: BLE001 # cache and codec failures must not abort token resolution
verbose_logger.debug("MCP per-user token cache set failed (ignored): %s", exc)
@with_service_target(MCP_OAUTH_TOKENS_TARGET)
async def delete(self, user_id: str, server_id: str) -> None:
await self.evict_token(user_id, server_id)
async def evict_token(self, user_id: str, server_id: str) -> None:
try:
await self.cache.async_delete_cache(self._key(user_id, server_id))
except Exception as exc: # noqa: BLE001
self.cache.in_memory_cache.delete_cache(self._key(user_id, server_id))
store: Final = self._store()
if store is not None:
await store.delete(user_id, server_id)
except Exception as exc: # noqa: BLE001 # best-effort eviction must not abort token resolution
verbose_logger.debug("MCP per-user token cache delete failed (ignored): %s", exc)

58
litellm/proxy/mcp_server/outbound_credentials/token_cache_codec.py
View file

@ -1,57 +1,3 @@
"""Serialize + encrypt boundary for caching an OAuth token in a shared (Redis) cache.
Shared cache values contain an encrypted access token and optional identity-binding proof.
Refresh tokens remain in the database; cache TTL bounds the access token's lifetime.
Legacy bearer-only entries decode without proof and cannot satisfy identity enforcement.
"""
from __future__ import annotations
import json
from collections.abc import Callable
from dataclasses import dataclass
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import (
OAuthToken,
from litellm.redis.stores.mcp_per_user_tokens import (
OAuthTokenCacheCodec as OAuthTokenCacheCodec, # noqa: PLC0414 # public compatibility export
)
_BOUND_PREFIX: Final = "litellm-bound-oauth-v1:"
_BOUND_PAYLOAD: Final = TypeAdapter(dict[str, str])
@dataclass(frozen=True, slots=True)
class OAuthTokenCacheCodec:
encrypt: Callable[[str], str]
decrypt: Callable[[str], str | None]
def encode(self, token: OAuthToken) -> str:
if token.identity_binding_proof is not None:
return self.encrypt(
_BOUND_PREFIX
+ json.dumps(
{
"access_token": token.access_token,
"identity_binding_proof": token.identity_binding_proof,
}
)
)
return self.encrypt(token.access_token)
def decode(self, blob: str) -> OAuthToken | None:
access_token: Final = self.decrypt(blob)
if not access_token:
return None
if access_token.startswith(_BOUND_PREFIX):
try:
payload: Final = _BOUND_PAYLOAD.validate_json(access_token[len(_BOUND_PREFIX) :])
except ValidationError:
return None
bearer: Final = payload.get("access_token")
proof: Final = payload.get("identity_binding_proof")
if not bearer or not proof:
return None
return OAuthToken(access_token=bearer, identity_binding_proof=proof)
return OAuthToken(access_token=access_token, refresh_token=None)

96
litellm/redis/stores/mcp_per_user_tokens.py Normal file
View file

@ -0,0 +1,96 @@
import ast
import json
from collections.abc import Callable
from dataclasses import dataclass
from datetime import timedelta
from typing import Final, Protocol
from pydantic import TypeAdapter, ValidationError
from litellm.constants import MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX
from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
from litellm.redis.transport.contract import RedisFailure, RedisTransport
TOKEN_KEY_PREFIX: Final = MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX + ":"
_BOUND_PREFIX: Final = "litellm-bound-oauth-v1:"
_BOUND_PAYLOAD: Final = TypeAdapter(dict[str, str])
_STRING: Final = TypeAdapter(str)
def per_user_token_key(user_id: str, server_id: str, *, prefix: str = TOKEN_KEY_PREFIX) -> str:
return f"{prefix}{user_id}:{server_id}"
@dataclass(frozen=True, slots=True)
class OAuthTokenCacheCodec:
encrypt: Callable[[str], str]
decrypt: Callable[[str], str | None]
def encode(self, token: OAuthToken) -> str:
if token.identity_binding_proof is not None:
return self.encrypt(
_BOUND_PREFIX
+ json.dumps(
{"access_token": token.access_token, "identity_binding_proof": token.identity_binding_proof}
)
)
return self.encrypt(token.access_token)
def decode(self, blob: str) -> OAuthToken | None:
access_token: Final = self.decrypt(blob)
if not access_token:
return None
if access_token.startswith(_BOUND_PREFIX):
try:
payload: Final = _BOUND_PAYLOAD.validate_json(access_token[len(_BOUND_PREFIX) :])
except ValidationError:
return None
bearer: Final = payload.get("access_token")
proof: Final = payload.get("identity_binding_proof")
if not bearer or not proof:
return None
return OAuthToken(access_token=bearer, identity_binding_proof=proof)
return OAuthToken(access_token=access_token, refresh_token=None)
class McpPerUserTokenStore(Protocol):
async def read(self, user_id: str, server_id: str) -> str | None | RedisFailure: ...
async def write(
self, user_id: str, server_id: str, ciphertext: str, *, ttl: timedelta, defer: bool = False
) -> bool | RedisFailure: ...
async def delete(self, user_id: str, server_id: str) -> int | RedisFailure: ...
@dataclass(frozen=True, slots=True)
class TransportMcpPerUserTokenStore:
transport: RedisTransport
prefix: str = TOKEN_KEY_PREFIX
async def read(self, user_id: str, server_id: str) -> str | None | RedisFailure:
reply: Final = await self.transport.get(per_user_token_key(user_id, server_id, prefix=self.prefix))
if reply is None or isinstance(reply, RedisFailure):
return reply
try:
text: Final = reply.decode("utf-8")
try:
return _STRING.validate_json(text)
except ValidationError:
return _STRING.validate_python(ast.literal_eval(text))
except (ValueError, SyntaxError):
return RedisFailure(kind="decode", message="Invalid MCP per-user token cache value")
async def write(
self, user_id: str, server_id: str, ciphertext: str, *, ttl: timedelta, defer: bool = False
) -> bool | RedisFailure:
key: Final = per_user_token_key(user_id, server_id, prefix=self.prefix)
wire: Final = json.dumps(ciphertext)
batch: Final = self.transport.request_batch() if defer else None
if batch is not None:
batch.set_seconds(key, wire, ttl=ttl)
return True
return await self.transport.set_seconds(key, wire, ttl=ttl)
async def delete(self, user_id: str, server_id: str) -> int | RedisFailure:
return await self.transport.delete(per_user_token_key(user_id, server_id, prefix=self.prefix))

95
tests/unit/proxy/mcp_server/outbound_credentials/test_dual_cache_token_backend.py
View file

@ -1,7 +1,11 @@
"""Tests for the DualCache-backed token cache backend: encrypted round-trip, key, TTL, miss."""
from typing import Final
import fakeredis
import pytest
from litellm.caching.dual_cache import DualCache
from litellm.proxy.mcp_server.outbound_credentials.dual_cache_token_backend import (
DualCacheTokenCacheBackend,
)
@ -11,25 +15,21 @@ from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import (
from litellm.proxy.mcp_server.outbound_credentials.token_cache_codec import (
OAuthTokenCacheCodec,
)
from tests._support.fake_redis import fake_redis_cache
class _FakeCache:
def __init__(self):
self.values = {}
self.ttls = {}
async def async_get_cache(self, key):
return self.values.get(key)
async def async_set_cache(self, key, value, ttl=None):
self.values[key] = value
self.ttls[key] = ttl
async def async_delete_cache(self, key):
self.values.pop(key, None)
def _cache() -> DualCache:
return DualCache(redis_cache=fake_redis_cache(fakeredis.FakeServer()))
def _backend(cache):
def _raw(cache: DualCache) -> fakeredis.FakeAsyncRedis:
assert cache.redis_cache is not None
client: Final = cache.redis_cache.init_async_client()
assert isinstance(client, fakeredis.FakeAsyncRedis)
return client
def _backend(cache: DualCache) -> DualCacheTokenCacheBackend:
codec = OAuthTokenCacheCodec(
encrypt=lambda s: f"enc:{s}",
decrypt=lambda b: b[4:] if b.startswith("enc:") else None,
@ -38,16 +38,17 @@ def _backend(cache):
@pytest.mark.asyncio
async def test_set_then_get_round_trips_encrypted_under_the_per_user_key():
cache = _FakeCache()
async def test_set_then_get_round_trips_encrypted_under_the_per_user_key(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("time.time", lambda: 1_800_000_000.0)
cache = _cache()
backend = _backend(cache)
await backend.set("alice", "srv", OAuthToken(access_token="at"), 120.0)
key = "mcp:per_user_token:alice:srv"
assert cache.ttls[key] == 120.0
assert await _raw(cache).pttl(key) == 120000
# Stored via the codec (encrypted), not the bare token; the codec's own test proves real
# NaCl output hides the secret - here the fake encrypt just wraps, so we check it was applied.
assert cache.values[key] == "enc:at"
assert await _raw(cache).get(key) == b'"enc:at"'
got = await backend.get("alice", "srv")
assert got is not None and got.access_token == "at"
@ -55,13 +56,13 @@ async def test_set_then_get_round_trips_encrypted_under_the_per_user_key():
@pytest.mark.asyncio
async def test_get_missing_key_is_none():
assert await _backend(_FakeCache()).get("alice", "srv") is None
assert await _backend(_cache()).get("alice", "srv") is None
@pytest.mark.asyncio
async def test_non_str_cache_value_is_a_miss():
cache = _FakeCache()
cache.values["mcp:per_user_token:alice:srv"] = 12345 # corrupt / wrong type
cache = _cache()
cache.in_memory_cache.set_cache("mcp:per_user_token:alice:srv", 12345)
assert await _backend(cache).get("alice", "srv") is None
@ -69,21 +70,21 @@ async def test_non_str_cache_value_is_a_miss():
async def test_undecryptable_blob_is_a_miss():
# e.g. a master-key rotation leaves an entry the codec can't decrypt; it must read as a miss so
# the store re-reads the DB, not raise (the decrypt helper here returns None for unknown blobs).
cache = _FakeCache()
cache.values["mcp:per_user_token:alice:srv"] = "not-our-ciphertext"
cache = _cache()
cache.in_memory_cache.set_cache("mcp:per_user_token:alice:srv", "not-our-ciphertext")
assert await _backend(cache).get("alice", "srv") is None
@pytest.mark.asyncio
async def test_non_positive_ttl_is_not_written():
cache = _FakeCache()
cache = _cache()
await _backend(cache).set("alice", "srv", OAuthToken(access_token="at"), 0.0)
assert cache.values == {} # an already-expired token is not cached
assert await cache.async_get_cache("mcp:per_user_token:alice:srv") is None
@pytest.mark.asyncio
async def test_delete_removes_the_entry():
cache = _FakeCache()
cache = _cache()
backend = _backend(cache)
await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0)
await backend.delete("alice", "srv")
@ -92,7 +93,7 @@ async def test_delete_removes_the_entry():
@pytest.mark.asyncio
async def test_keys_isolate_users_and_servers():
cache = _FakeCache()
cache = _cache()
backend = _backend(cache)
await backend.set("alice", "srv", OAuthToken(access_token="a"), 60.0)
await backend.set("bob", "srv", OAuthToken(access_token="b"), 60.0)
@ -102,35 +103,30 @@ async def test_keys_isolate_users_and_servers():
assert bob is not None and bob.access_token == "b"
class _RaisingCache(_FakeCache):
"""A cache whose every op raises, e.g. a Redis outage tripping the circuit breaker."""
async def async_get_cache(self, key):
raise ConnectionError("redis down")
async def async_set_cache(self, key, value, ttl=None):
raise ConnectionError("redis down")
async def async_delete_cache(self, key):
raise ConnectionError("redis down")
def _offline_cache():
server = fakeredis.FakeServer()
server.connected = False
return DualCache(redis_cache=fake_redis_cache(server))
# A cache or codec failure must degrade to the safe value (miss / no-op), never propagate: otherwise a
# Redis outage turns CachedOAuthTokenStore.fetch() into a 500 instead of a cache miss that re-reads the
# DB (get/set) or issues the OAuth challenge (the unauthorized branch deletes before returning None).
@pytest.mark.asyncio
async def test_get_is_a_miss_when_the_cache_raises():
assert await _backend(_RaisingCache()).get("alice", "srv") is None
assert await _backend(_offline_cache()).get("alice", "srv") is None
@pytest.mark.asyncio
async def test_set_is_swallowed_when_the_cache_raises():
await _backend(_RaisingCache()).set("alice", "srv", OAuthToken(access_token="at"), 60.0)
backend = _backend(_offline_cache())
await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0)
assert await backend.get("alice", "srv") == OAuthToken(access_token="at")
@pytest.mark.asyncio
async def test_delete_is_swallowed_when_the_cache_raises():
await _backend(_RaisingCache()).delete("alice", "srv")
backend = _backend(_offline_cache())
await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0)
await backend.delete("alice", "srv")
assert await backend.get("alice", "srv") is None
@pytest.mark.asyncio
@ -139,8 +135,9 @@ async def test_set_is_swallowed_when_the_codec_raises():
raise ValueError("encrypt unavailable")
codec = OAuthTokenCacheCodec(encrypt=_boom, decrypt=lambda b: None)
backend = DualCacheTokenCacheBackend(_FakeCache(), codec)
await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0) # must not raise
backend = DualCacheTokenCacheBackend(_cache(), codec)
await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0)
assert await backend.get("alice", "srv") is None
@pytest.mark.asyncio
@ -151,6 +148,6 @@ async def test_get_is_a_miss_when_decrypt_raises():
raise ValueError("bad ciphertext")
codec = OAuthTokenCacheCodec(encrypt=lambda s: f"enc:{s}", decrypt=_boom)
cache = _FakeCache()
cache.values["mcp:per_user_token:alice:srv"] = "whatever"
cache = _cache()
cache.in_memory_cache.set_cache("mcp:per_user_token:alice:srv", "whatever")
assert await DualCacheTokenCacheBackend(cache, codec).get("alice", "srv") is None

70
tests/unit/proxy/mcp_server/test_per_user_oauth_cache.py
View file

@ -7,12 +7,15 @@ Unit tests for per-user MCP OAuth token storage:
"""
import sys
import time
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Optional
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.caching.dual_cache import DualCache
# Stub out modules that aren't available in the unit-test environment
# so we can import the targets without a full proxy stack.
for _mod in ("orjson",):
@ -22,7 +25,6 @@ for _mod in ("orjson",):
from litellm.proxy.mcp_server.oauth2_token_cache import ( # noqa: E402
MCPPerUserTokenCache,
compute_per_user_token_ttl,
mcp_per_user_token_cache,
)
from litellm.types.mcp import MCPAuth, MCPTransport # noqa: E402
from litellm.types.mcp_server.mcp_server_manager import MCPServer # noqa: E402
@ -41,7 +43,7 @@ def _import_validate():
def _make_server(**kwargs) -> MCPServer:
defaults: Dict[str, Any] = {
defaults: dict[str, Any] = {
"server_id": "slack-test",
"name": "Slack",
"server_name": "slack",
@ -262,9 +264,7 @@ class TestMCPPerUserTokenCache:
@pytest.mark.asyncio
async def test_get_returns_none_on_miss(self, cache, mock_dual_cache):
with (
patch(
"litellm.proxy.mcp_server.oauth2_token_cache.decrypt_value_helper"
) as mock_decrypt,
patch("litellm.proxy.mcp_server.oauth2_token_cache.decrypt_value_helper") as mock_decrypt,
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache),
):
mock_dual_cache.async_get_cache.return_value = None
@ -273,66 +273,46 @@ class TestMCPPerUserTokenCache:
mock_decrypt.assert_not_called()
@pytest.mark.asyncio
async def test_get_decrypts_cached_value(self, cache, mock_dual_cache):
fake_encrypted = "encrypted_blob_abc123"
fake_plaintext = "xoxb-slack-token"
async def test_get_decrypts_cached_value(self, cache: MCPPerUserTokenCache) -> None:
dual: Final = DualCache()
await dual.async_set_cache("mcp:per_user_token:alice:slack-test", "encrypted_blob_xyz", ttl=60)
with (
patch(
"litellm.proxy.mcp_server.oauth2_token_cache.decrypt_value_helper",
return_value=fake_plaintext,
return_value="xoxb-slack-token",
) as mock_decrypt,
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache),
patch("litellm.proxy.proxy_server.user_api_key_cache", dual),
):
mock_dual_cache.async_get_cache.return_value = fake_encrypted
result = await cache.get("alice", "slack-test")
assert result == fake_plaintext
mock_decrypt.assert_called_once_with(
fake_encrypted,
key="mcp_per_user_token",
exception_type="debug",
)
assert result == "xoxb-slack-token"
mock_decrypt.assert_called_once_with("encrypted_blob_xyz", key="mcp_per_user_token", exception_type="debug")
@pytest.mark.asyncio
async def test_set_encrypts_before_storing(self, cache, mock_dual_cache):
fake_encrypted = "encrypted_blob_xyz"
async def test_set_encrypts_before_storing_under_the_exact_key(self, cache: MCPPerUserTokenCache) -> None:
dual: Final = DualCache()
with (
patch(
"litellm.proxy.mcp_server.oauth2_token_cache.encrypt_value_helper",
return_value=fake_encrypted,
return_value="encrypted_blob_xyz",
) as mock_encrypt,
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache),
patch("litellm.proxy.proxy_server.user_api_key_cache", dual),
):
await cache.set("alice", "slack-test", "xoxb-token", ttl=3540)
mock_encrypt.assert_called_once_with("xoxb-token")
mock_dual_cache.async_set_cache.assert_called_once()
call_kwargs = mock_dual_cache.async_set_cache.call_args
assert call_kwargs[0][1] == fake_encrypted # encrypted value stored
assert call_kwargs[1]["ttl"] == 3540
assert dual.in_memory_cache.get_cache("mcp:per_user_token:alice:slack-test") == "encrypted_blob_xyz"
assert dual.in_memory_cache.get_cache("xoxb-token") is None
assert 0 < dual.in_memory_cache.ttl_dict["mcp:per_user_token:alice:slack-test"] - time.time() <= 3540
@pytest.mark.asyncio
async def test_set_uses_correct_cache_key(self, cache, mock_dual_cache):
with (
patch(
"litellm.proxy.mcp_server.oauth2_token_cache.encrypt_value_helper",
return_value="enc",
),
patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache),
):
await cache.set("bob", "github-server", "ghp_token", ttl=3600)
key_used = mock_dual_cache.async_set_cache.call_args[0][0]
assert key_used == "mcp:per_user_token:bob:github-server"
@pytest.mark.asyncio
async def test_delete_calls_async_delete_cache(self, cache, mock_dual_cache):
mock_dual_cache.async_delete_cache = AsyncMock()
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache):
async def test_delete_removes_the_cached_token(self, cache: MCPPerUserTokenCache) -> None:
dual: Final = DualCache()
await dual.async_set_cache("mcp:per_user_token:alice:slack-test", "encrypted_blob_xyz", ttl=60)
with patch("litellm.proxy.proxy_server.user_api_key_cache", dual):
await cache.delete("alice", "slack-test")
mock_dual_cache.async_delete_cache.assert_called_once_with(key="mcp:per_user_token:alice:slack-test")
mock_dual_cache.async_set_cache.assert_not_called()
assert dual.in_memory_cache.get_cache("mcp:per_user_token:alice:slack-test") is None
@pytest.mark.asyncio
async def test_get_returns_none_on_decrypt_failure(self, cache, mock_dual_cache):

224
tests/unit/redis/stores/test_mcp_per_user_tokens.py Normal file
View file

@ -0,0 +1,224 @@
import json
from contextlib import nullcontext
from datetime import timedelta
from typing import Final
from unittest.mock import patch
import fakeredis
import pytest
from pydantic import TypeAdapter
from redis.exceptions import ConnectionError
from litellm.caching.redis_batch import active_request_redis_batch, request_redis_batch_scope
from litellm.proxy import proxy_server
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
AuthCacheInvalidationSubscriber,
auth_cache_invalidation_channel,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.mcp_server.oauth2_token_cache import MCPPerUserTokenCache
from litellm.proxy.mcp_server.outbound_credentials.dual_cache_token_backend import DualCacheTokenCacheBackend
from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
from litellm.proxy.mcp_server.outbound_credentials.token_cache_codec import OAuthTokenCacheCodec
from litellm.redis.stores.mcp_per_user_tokens import TransportMcpPerUserTokenStore
from litellm.redis.transport.contract import RedisFailure
from tests._support.fake_redis import fake_redis_cache, trip_breaker
@pytest.fixture(autouse=True)
def encryption_and_clock(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(proxy_server, "master_key", "sk-test-master")
monkeypatch.setattr("time.time", lambda: 1_800_000_000.0)
_MESSAGE: Final = TypeAdapter(dict[str, object])
def _legacy_codec(cache: MCPPerUserTokenCache) -> OAuthTokenCacheCodec:
return cache._codec() # pyright: ignore[reportPrivateUsage] # the legacy cache owns the codec the new backend must match
@pytest.mark.asyncio
@pytest.mark.parametrize("proof", [None, "bound-é"])
async def test_legacy_wire_is_shared_by_both_callers_with_exact_codec_and_expiry(
monkeypatch: pytest.MonkeyPatch,
proof: str | None,
) -> None:
server: Final = fakeredis.FakeServer()
redis: Final = fake_redis_cache(server, namespace="oauth")
raw: Final = fakeredis.FakeAsyncRedis(server=server)
cache: Final = UserApiKeyCache(redis_cache=redis)
legacy: Final = MCPPerUserTokenCache()
backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(legacy))
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
user: Final = 'a:é"'
key: Final = f"mcp:per_user_token:{user}:東京"
token: Final = OAuthToken("bearer-é", identity_binding_proof=proof)
plaintext: Final = (
token.access_token
if proof is None
else "litellm-bound-oauth-v1:"
+ json.dumps({"access_token": token.access_token, "identity_binding_proof": proof})
)
ciphertext: Final = encrypt_value_helper(plaintext)
await redis.async_set_cache(key, ciphertext, ttl=90)
assert await backend.get(user, "東京") == token
assert await legacy.get_token(user, "東京") == token
await backend.set(user, "東京", token, 120.9)
written: Final = await redis.async_get_cache(key)
assert isinstance(written, str)
assert decrypt_value_helper(written, "test") == plaintext
assert await raw.get("oauth:" + key) == json.dumps(written).encode()
assert await raw.pttl("oauth:" + key) == 120000
assert await raw.keys("*") == [("oauth:" + key).encode()]
cache.in_memory_cache.delete_cache(key)
await legacy.set(user, "東京", "updated", 73, proof)
cache.in_memory_cache.delete_cache(key)
assert await backend.get(user, "東京") == OAuthToken("updated", identity_binding_proof=proof)
assert await raw.pttl("oauth:" + key) == 73000
await backend.delete(user, "東京")
assert await legacy.get(user, "東京") is None
assert await raw.get("oauth:" + key) is None
@pytest.mark.asyncio
async def test_auth_cache_deferred_write_joins_request_batch_and_retains_memory_policy() -> None:
server: Final = fakeredis.FakeServer()
redis: Final = fake_redis_cache(server, namespace="oauth")
raw: Final = fakeredis.FakeAsyncRedis(server=server)
cache: Final = UserApiKeyCache(redis_cache=redis)
backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(MCPPerUserTokenCache()))
with request_redis_batch_scope():
await backend.set("u", "s", OAuthToken("pending"), 20.9)
assert await backend.get("u", "s") == OAuthToken("pending")
assert await raw.get("oauth:mcp:per_user_token:u:s") is None
batch: Final = active_request_redis_batch(redis)
assert batch is not None
await batch.flush()
assert await raw.pttl("oauth:mcp:per_user_token:u:s") == 20000
cache.in_memory_cache.delete_cache("mcp:per_user_token:u:s")
assert await backend.get("u", "s") == OAuthToken("pending")
batch.delete("mcp:per_user_token:u:s")
cache.in_memory_cache.delete_cache("mcp:per_user_token:u:s")
assert await backend.get("u", "s") is None
await batch.flush()
assert await raw.get("oauth:mcp:per_user_token:u:s") is None
@pytest.mark.asyncio
@pytest.mark.parametrize("wire", [b"{}", b"null", b"42", b"broken", b"\xff", b'"litellm-bound-oauth-v1:{}"'])
async def test_malformed_redis_tokens_are_safe_misses(monkeypatch: pytest.MonkeyPatch, wire: bytes) -> None:
server: Final = fakeredis.FakeServer()
redis: Final = fake_redis_cache(server)
cache: Final = UserApiKeyCache(redis_cache=redis)
legacy: Final = MCPPerUserTokenCache()
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
await fakeredis.FakeAsyncRedis(server=server).set("mcp:per_user_token:u:s", wire)
assert await legacy.get("u", "s") is None
assert await DualCacheTokenCacheBackend(cache, _legacy_codec(legacy)).get("u", "s") is None
@pytest.mark.asyncio
@pytest.mark.parametrize("applied", [False, True])
async def test_ambiguous_write_retains_local_value_without_retry_and_peer_sees_only_applied_value(
applied: bool,
) -> None:
server: Final = fakeredis.FakeServer()
redis: Final = fake_redis_cache(server, namespace="oauth")
raw: Final = fakeredis.FakeAsyncRedis(server=server)
local: Final = DualCacheTokenCacheBackend(UserApiKeyCache(redis_cache=redis), _legacy_codec(MCPPerUserTokenCache()))
peer: Final = DualCacheTokenCacheBackend(UserApiKeyCache(redis_cache=redis), _legacy_codec(MCPPerUserTokenCache()))
async def lost_reply(*args: object, **kwargs: object) -> None:
if applied:
assert isinstance(args[1], (str, bytes))
await raw.set("oauth:mcp:per_user_token:u:s", args[1], ex=20)
raise ConnectionError("lost reply")
with patch.object(redis.init_async_client(), "set", lost_reply):
await local.set("u", "s", OAuthToken("test-token"), 20)
assert await local.get("u", "s") == OAuthToken("test-token")
assert await peer.get("u", "s") == (OAuthToken("test-token") if applied else None)
@pytest.mark.asyncio
async def test_expiry_and_outage_preserve_local_cache_and_safe_misses(monkeypatch: pytest.MonkeyPatch) -> None:
server: Final = fakeredis.FakeServer()
redis: Final = fake_redis_cache(server)
cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=5)
backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(MCPPerUserTokenCache()))
await backend.set("u", "s", OAuthToken("expires"), 2)
monkeypatch.setattr("time.time", lambda: 1_800_000_003.0)
assert await backend.get("u", "s") is None
server.connected = False
await backend.set("u", "s", OAuthToken("local"), 20)
assert await backend.get("u", "s") == OAuthToken("local")
await backend.delete("u", "s")
assert await backend.get("u", "s") is None
store: Final = TransportMcpPerUserTokenStore(redis.transport())
assert isinstance(await store.read("u", "s"), RedisFailure)
assert isinstance(await store.write("u", "s", "encrypted", ttl=timedelta(seconds=2)), RedisFailure)
@pytest.mark.asyncio
@pytest.mark.parametrize("deferred", [False, True])
@pytest.mark.parametrize("ttl", [0.9, 0.99999999])
async def test_subsecond_ttl_keeps_local_value_but_does_not_replace_shared_entry(deferred: bool, ttl: float) -> None:
server: Final = fakeredis.FakeServer()
redis: Final = fake_redis_cache(server)
cache: Final = UserApiKeyCache(redis_cache=redis)
backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(MCPPerUserTokenCache()))
await backend.set("u", "s", OAuthToken("old"), 30)
with request_redis_batch_scope() if deferred else nullcontext():
await backend.set("u", "s", OAuthToken("local"), ttl)
assert await backend.get("u", "s") == OAuthToken("local")
batch: Final = active_request_redis_batch(redis)
if batch is not None:
await batch.flush()
cache.in_memory_cache.delete_cache("mcp:per_user_token:u:s")
assert await backend.get("u", "s") == OAuthToken("old")
@pytest.mark.asyncio
async def test_open_breaker_preserves_existing_auth_memory_fallback() -> None:
redis: Final = fake_redis_cache(fakeredis.FakeServer())
trip_breaker(redis)
store: Final = TransportMcpPerUserTokenStore(redis.transport())
read: Final = await store.read("u", "s")
assert isinstance(read, RedisFailure) and read.kind == "breaker_open"
cache: Final = UserApiKeyCache(redis_cache=redis)
backend: Final = DualCacheTokenCacheBackend(cache, _legacy_codec(MCPPerUserTokenCache()))
await backend.set("u", "s", OAuthToken("local"), 60)
assert await backend.get("u", "s") == OAuthToken("local")
await backend.delete("u", "s")
assert await backend.get("u", "s") is None
@pytest.mark.asyncio
async def test_legacy_delete_evicts_redis_local_and_peer_through_auth_invalidation(
monkeypatch: pytest.MonkeyPatch,
) -> None:
server: Final = fakeredis.FakeServer()
redis: Final = fake_redis_cache(server, namespace="oauth")
local: Final = UserApiKeyCache(redis_cache=redis)
peer: Final = UserApiKeyCache(redis_cache=redis)
legacy: Final = MCPPerUserTokenCache()
backend: Final = DualCacheTokenCacheBackend(peer, _legacy_codec(legacy))
monkeypatch.setattr(proxy_server, "user_api_key_cache", local)
monkeypatch.setattr(proxy_server, "redis_usage_cache", redis)
await legacy.set("u", "s", "revoke-me", 60)
assert await backend.get("u", "s") == OAuthToken("revoke-me")
pubsub: Final = fakeredis.FakeAsyncRedis(server=server).pubsub()
await pubsub.subscribe(auth_cache_invalidation_channel(redis))
await pubsub.get_message()
await legacy.delete("u", "s")
message: Final = _MESSAGE.validate_python(await pubsub.get_message(ignore_subscribe_messages=True, timeout=1)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # redis-py pubsub replies are untyped
data: Final = message["data"]
assert isinstance(data, (str, bytes))
assert json.loads(data) == {"cache_key": "mcp:per_user_token:u:s"}
AuthCacheInvalidationSubscriber(redis, peer)._apply_message(message) # pyright: ignore[reportPrivateUsage] # drives the subscriber without a live pubsub loop
assert await backend.get("u", "s") is None
assert await legacy.get("u", "s") is None
assert await fakeredis.FakeAsyncRedis(server=server).get("oauth:mcp:per_user_token:u:s") is None
await pubsub.aclose() # pyright: ignore[reportAttributeAccessIssue] # redis-py stubs predate PubSub.aclose