diff --git a/litellm/proxy/mcp_server/oauth2_token_cache.py b/litellm/proxy/mcp_server/oauth2_token_cache.py index 1e6e83c0f4b..45d4efa7a8e 100644 --- a/litellm/proxy/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/mcp_server/oauth2_token_cache.py @@ -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", diff --git a/litellm/proxy/mcp_server/outbound_credentials/dual_cache_token_backend.py b/litellm/proxy/mcp_server/outbound_credentials/dual_cache_token_backend.py index b360b787962..64bd0580b3c 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/dual_cache_token_backend.py +++ b/litellm/proxy/mcp_server/outbound_credentials/dual_cache_token_backend.py @@ -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) diff --git a/litellm/proxy/mcp_server/outbound_credentials/token_cache_codec.py b/litellm/proxy/mcp_server/outbound_credentials/token_cache_codec.py index 4576e925939..5f979da6bed 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/token_cache_codec.py +++ b/litellm/proxy/mcp_server/outbound_credentials/token_cache_codec.py @@ -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) diff --git a/litellm/redis/stores/mcp_per_user_tokens.py b/litellm/redis/stores/mcp_per_user_tokens.py new file mode 100644 index 00000000000..f7bf8cf17a7 --- /dev/null +++ b/litellm/redis/stores/mcp_per_user_tokens.py @@ -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)) diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/test_dual_cache_token_backend.py b/tests/unit/proxy/mcp_server/outbound_credentials/test_dual_cache_token_backend.py index 00db27a60b8..a298214e688 100644 --- a/tests/unit/proxy/mcp_server/outbound_credentials/test_dual_cache_token_backend.py +++ b/tests/unit/proxy/mcp_server/outbound_credentials/test_dual_cache_token_backend.py @@ -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 diff --git a/tests/unit/proxy/mcp_server/test_per_user_oauth_cache.py b/tests/unit/proxy/mcp_server/test_per_user_oauth_cache.py index c0e23319566..a953910c9be 100644 --- a/tests/unit/proxy/mcp_server/test_per_user_oauth_cache.py +++ b/tests/unit/proxy/mcp_server/test_per_user_oauth_cache.py @@ -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): diff --git a/tests/unit/redis/stores/test_mcp_per_user_tokens.py b/tests/unit/redis/stores/test_mcp_per_user_tokens.py new file mode 100644 index 00000000000..84f610c3433 --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_per_user_tokens.py @@ -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