diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 5e3a3946604..1da141ec63b 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -22,6 +22,11 @@ from litellm.redis.stores.mcp_agent_permission_links import ( TransportMcpAgentPermissionLinkStore, agent_permission_link_key, ) +from litellm.redis.stores.mcp_per_user_token import ( + McpPerUserTokenStore, + TransportMcpPerUserTokenStore, + per_user_token_key, +) from litellm.redis.stores.mcp_user_permission_links import ( USER_NO_MCP_PERMISSION_SENTINEL as _USER_NO_PERMISSION, ) @@ -412,6 +417,42 @@ async def fill_agent_permission_link( _ = await store.fill(agent_id, permission_id, ttl=timedelta(seconds=int(ttl))) +async def get_mcp_per_user_token(cache: UserApiKeyCache, user_id: str, server_id: str) -> str | None: + key: Final = per_user_token_key(user_id, server_id) + local: Final[object] = await cache.async_get_cache(key=key, local_only=True) + if local is not None: + return local if isinstance(local, str) else None + if cache.redis_cache is None: + return None + batch: Final = active_request_redis_batch(cache.redis_cache) + if batch is not None and batch.read_as_missing(key): + return None + store: Final[McpPerUserTokenStore] = TransportMcpPerUserTokenStore(cache.redis_cache.transport()) + blob: Final = await store.get(user_id, server_id) + if blob is None or isinstance(blob, RedisFailure): + return None + await cache.async_set_cache(key=key, value=blob, local_only=True) + return blob + + +async def set_mcp_per_user_token( + cache: UserApiKeyCache, user_id: str, server_id: str, blob: str, *, ttl: float +) -> None: + await cache.async_set_cache(key=per_user_token_key(user_id, server_id), value=blob, ttl=ttl, local_only=True) + if cache.redis_cache is None: + return + store: Final[McpPerUserTokenStore] = TransportMcpPerUserTokenStore(cache.redis_cache.transport()) + _ = await store.put(user_id, server_id, blob, ttl=timedelta(seconds=int(ttl))) + + +async def delete_mcp_per_user_token(cache: UserApiKeyCache, user_id: str, server_id: str) -> None: + cache.in_memory_cache.delete_cache(per_user_token_key(user_id, server_id)) + if cache.redis_cache is None: + return + store: Final[McpPerUserTokenStore] = TransportMcpPerUserTokenStore(cache.redis_cache.transport()) + _ = await store.delete(user_id, server_id) + + #: Cached under ``tag_registry_cache_key`` when the table exceeds ``TAG_REGISTRY_MAX_SIZE``: #: registry unusable, fall back to the per-tag lookup. TAG_REGISTRY_OVERFLOW_SENTINEL: Final = "__tag_registry_overflow__" diff --git a/litellm/proxy/mcp_server/oauth2_token_cache.py b/litellm/proxy/mcp_server/oauth2_token_cache.py index 1e6e83c0f4b..694fa9d55ce 100644 --- a/litellm/proxy/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/mcp_server/oauth2_token_cache.py @@ -23,7 +23,6 @@ from litellm.constants import ( 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 ( @@ -37,6 +36,7 @@ from litellm.proxy.mcp_server.oauth_utils import ( 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_token import per_user_token_key from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -241,7 +241,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( @@ -256,10 +256,10 @@ class MCPPerUserTokenCache: @with_service_target(MCP_OAUTH_TOKENS_TARGET) async def get_token(self, user_id: str, server_id: str) -> OAuthToken | None: try: + from litellm.proxy.common_utils.user_api_key_cache import get_mcp_per_user_token # noqa: PLC0415 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) + encrypted: Final = await get_mcp_per_user_token(user_api_key_cache, user_id, server_id) if encrypted is None: return None return self._codec().decode(encrypted) @@ -283,13 +283,13 @@ class MCPPerUserTokenCache: ) -> None: """Store NaCl-encrypted access_token in Redis with the given TTL.""" try: + from litellm.proxy.common_utils.user_api_key_cache import set_mcp_per_user_token # noqa: PLC0415 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 user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl) + await set_mcp_per_user_token(user_api_key_cache, user_id, server_id, encrypted, ttl=ttl) verbose_logger.debug( "MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds", user_id, @@ -308,12 +308,13 @@ class MCPPerUserTokenCache: """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.common_utils.user_api_key_cache import delete_mcp_per_user_token # noqa: PLC0415 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 delete_mcp_per_user_token(user_api_key_cache, user_id, server_id) + await publish_auth_cache_invalidation(cache_key=self._cache_key(user_id, server_id)) 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..babb0aa3e55 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 @@ -9,11 +9,17 @@ 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 typing import Final, Protocol +from dataclasses import dataclass +from typing import Final from litellm._internal_context import with_service_target from litellm._logging import verbose_logger +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + delete_mcp_per_user_token, + get_mcp_per_user_token, + set_mcp_per_user_token, +) from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, ) @@ -23,16 +29,6 @@ from litellm.proxy.mcp_server.outbound_credentials.token_cache_codec import ( from litellm.proxy.mcp_server.utils import MCP_OAUTH_TOKENS_TARGET -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: ... - - async def async_set_cache(self, key: str, value: str, ttl: float | None = None) -> None: ... - - async def async_delete_cache(self, key: str) -> None: ... - - @dataclass(frozen=True, slots=True) class DualCacheTokenCacheBackend: """Every method degrades a cache or codec failure to its safe value - ``get`` to a miss @@ -41,19 +37,14 @@ class DualCacheTokenCacheBackend: contract. The guarantee holds here regardless of whether the injected cache/codec also swallow. """ - cache: AsyncCache + cache: UserApiKeyCache codec: OAuthTokenCacheCodec - _: KW_ONLY - key_prefix: str = "mcp:per_user_token:" - - def _key(self, user_id: str, server_id: str) -> str: - return f"{self.key_prefix}{user_id}:{server_id}" @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)) - return self.codec.decode(blob) if isinstance(blob, str) else None + blob: Final = await get_mcp_per_user_token(self.cache, user_id, server_id) + return self.codec.decode(blob) if blob is not None else None except Exception as exc: # noqa: BLE001 verbose_logger.debug("MCP per-user token cache get failed (miss): %s", exc) return None @@ -63,17 +54,13 @@ class DualCacheTokenCacheBackend: if ttl_seconds <= 0: return try: - await self.cache.async_set_cache( - self._key(user_id, server_id), - self.codec.encode(token), - ttl=ttl_seconds, - ) + await set_mcp_per_user_token(self.cache, user_id, server_id, self.codec.encode(token), ttl=ttl_seconds) except Exception as exc: # noqa: BLE001 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: try: - await self.cache.async_delete_cache(self._key(user_id, server_id)) + await delete_mcp_per_user_token(self.cache, user_id, server_id) except Exception as exc: # noqa: BLE001 verbose_logger.debug("MCP per-user token cache delete failed (ignored): %s", exc) diff --git a/litellm/proxy/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/mcp_server/outbound_credentials/per_user_oauth_store.py index f3748f99ba4..7c3df58ee6d 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -21,7 +21,6 @@ from litellm.proxy.mcp_server.outbound_credentials.authz_code_refresher import ( AuthorizationCodeRefresher, ) from litellm.proxy.mcp_server.outbound_credentials.dual_cache_token_backend import ( - AsyncCache, DualCacheTokenCacheBackend, ) from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( @@ -141,10 +140,7 @@ def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, Refres encrypt_value_helper, lambda blob: decrypt_value_helper(blob, "mcp_per_user_token", exception_type="debug"), ) - # user_api_key_cache satisfies the AsyncCache slice (DualCache types ttl via **kwargs) - an - # untyped-boundary cast. - cache: Final[AsyncCache] = user_api_key_cache # pyright: ignore - backend: Final = DualCacheTokenCacheBackend(cache, codec) + backend: Final = DualCacheTokenCacheBackend(user_api_key_cache, codec) return backend, coordinator, True diff --git a/litellm/redis/stores/mcp_per_user_token.py b/litellm/redis/stores/mcp_per_user_token.py new file mode 100644 index 00000000000..441e16d77c5 --- /dev/null +++ b/litellm/redis/stores/mcp_per_user_token.py @@ -0,0 +1,54 @@ +import ast +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter + +from litellm.constants import MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_DECODED: Final = TypeAdapter(object) + + +def per_user_token_key(user_id: str, server_id: str) -> str: + return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}" + + +def _decode(raw: bytes) -> str | None | RedisFailure: + text: Final = raw.decode("utf-8", errors="replace") + try: + value: Final = _DECODED.validate_python(json.loads(text)) + except ValueError: + try: + value_literal: Final = _DECODED.validate_python(ast.literal_eval(text)) + except (ValueError, SyntaxError) as error: + return RedisFailure(kind="decode", message=str(error)) + return value_literal if isinstance(value_literal, str) else None + return value if isinstance(value, str) else None + + +class McpPerUserTokenStore(Protocol): + async def get(self, user_id: str, server_id: str) -> str | None | RedisFailure: ... + + async def put(self, user_id: str, server_id: str, blob: str, *, ttl: timedelta) -> bool | RedisFailure: ... + + async def delete(self, user_id: str, server_id: str) -> int | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportMcpPerUserTokenStore: + transport: RedisTransport + + async def get(self, user_id: str, server_id: str) -> str | None | RedisFailure: + raw: Final = await self.transport.get(per_user_token_key(user_id, server_id)) + if raw is None or isinstance(raw, RedisFailure): + return raw + return _decode(raw) + + async def put(self, user_id: str, server_id: str, blob: str, *, ttl: timedelta) -> bool | RedisFailure: + return await self.transport.set_seconds(per_user_token_key(user_id, server_id), json.dumps(blob), 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)) 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..86081679bb8 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,156 +1,114 @@ -"""Tests for the DualCache-backed token cache backend: encrypted round-trip, key, TTL, miss.""" +"""The shared per-user token backend: encrypted round trip, legacy wire bytes, TTL, misses and outages.""" +import json + +import fakeredis import pytest -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.caching.in_memory_cache import InMemoryCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +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 tests._support.fake_redis import fake_redis_cache + +KEY = "mcp:per_user_token:alice:srv" +CODEC = OAuthTokenCacheCodec(encrypt=lambda s: f"enc:{s}", decrypt=lambda b: b[4:] if b.startswith("enc:") else None) -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 _worker(server: fakeredis.FakeServer, namespace: str | None = None) -> UserApiKeyCache: + return UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=fake_redis_cache(server, namespace=namespace)) -def _backend(cache): - codec = OAuthTokenCacheCodec( - encrypt=lambda s: f"enc:{s}", - decrypt=lambda b: b[4:] if b.startswith("enc:") else None, - ) +def _backend(cache: UserApiKeyCache, codec: OAuthTokenCacheCodec = CODEC) -> DualCacheTokenCacheBackend: return DualCacheTokenCacheBackend(cache, codec) @pytest.mark.asyncio -async def test_set_then_get_round_trips_encrypted_under_the_per_user_key(): - cache = _FakeCache() - backend = _backend(cache) - await backend.set("alice", "srv", OAuthToken(access_token="at"), 120.0) +async def test_a_token_one_worker_writes_reaches_another_with_the_legacy_wire_bytes() -> None: + server = fakeredis.FakeServer() + raw = fakeredis.FakeAsyncRedis(server=server) + await _backend(_worker(server, "ns")).set("alice", "srv", OAuthToken(access_token="at"), 120.9) - key = "mcp:per_user_token:alice:srv" - assert cache.ttls[key] == 120.0 - # 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" - - got = await backend.get("alice", "srv") + assert await raw.get(f"ns:{KEY}") == json.dumps("enc:at").encode() + assert await raw.ttl(f"ns:{KEY}") == 120 + got = await _backend(_worker(server, "ns")).get("alice", "srv") assert got is not None and got.access_token == "at" @pytest.mark.asyncio -async def test_get_missing_key_is_none(): - assert await _backend(_FakeCache()).get("alice", "srv") is None +async def test_a_token_an_old_pod_wrote_is_read_and_cached_locally() -> None: + server = fakeredis.FakeServer() + raw = fakeredis.FakeAsyncRedis(server=server) + legacy = fake_redis_cache(server) + await legacy.async_set_cache(KEY, "enc:old", ttl=60) + reader = _worker(server) + + got = await _backend(reader).get("alice", "srv") + await raw.delete(KEY) + cached = await _backend(reader).get("alice", "srv") + + assert got is not None and got.access_token == "old" + assert cached is not None and cached.access_token == "old" @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 - assert await _backend(cache).get("alice", "srv") is None +@pytest.mark.parametrize("stored", [b"12345", b'"not-our-ciphertext"', b"{not json"]) +async def test_wrong_type_undecryptable_or_malformed_entries_are_misses(stored: bytes) -> None: + server = fakeredis.FakeServer() + await fakeredis.FakeAsyncRedis(server=server).set(KEY, stored) + assert await _backend(_worker(server)).get("alice", "srv") is None @pytest.mark.asyncio -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" - assert await _backend(cache).get("alice", "srv") is None +async def test_non_positive_ttl_is_not_written() -> None: + server = fakeredis.FakeServer() + await _backend(_worker(server)).set("alice", "srv", OAuthToken(access_token="at"), 0.0) + assert await fakeredis.FakeAsyncRedis(server=server).keys("*") == [] @pytest.mark.asyncio -async def test_non_positive_ttl_is_not_written(): - cache = _FakeCache() - await _backend(cache).set("alice", "srv", OAuthToken(access_token="at"), 0.0) - assert cache.values == {} # an already-expired token is not cached - - -@pytest.mark.asyncio -async def test_delete_removes_the_entry(): - cache = _FakeCache() - backend = _backend(cache) +async def test_delete_removes_the_entry_here_and_in_redis() -> None: + server = fakeredis.FakeServer() + backend = _backend(_worker(server)) await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0) await backend.delete("alice", "srv") assert await backend.get("alice", "srv") is None + assert await fakeredis.FakeAsyncRedis(server=server).exists(KEY) == 0 @pytest.mark.asyncio -async def test_keys_isolate_users_and_servers(): - cache = _FakeCache() - backend = _backend(cache) +async def test_keys_isolate_users_and_servers() -> None: + server = fakeredis.FakeServer() + backend = _backend(_worker(server)) await backend.set("alice", "srv", OAuthToken(access_token="a"), 60.0) await backend.set("bob", "srv", OAuthToken(access_token="b"), 60.0) - alice = await backend.get("alice", "srv") - bob = await backend.get("bob", "srv") + alice = await _backend(_worker(server)).get("alice", "srv") + bob = await _backend(_worker(server)).get("bob", "srv") assert alice is not None and alice.access_token == "a" 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") - - -# 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 +async def test_a_redis_outage_degrades_every_operation_to_its_safe_value() -> None: + server = fakeredis.FakeServer() + backend = _backend(_worker(server)) + server.connected = False + await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0) + await backend.delete("alice", "srv") + assert await _backend(_worker(server)).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) - - -@pytest.mark.asyncio -async def test_delete_is_swallowed_when_the_cache_raises(): - await _backend(_RaisingCache()).delete("alice", "srv") - - -@pytest.mark.asyncio -async def test_set_is_swallowed_when_the_codec_raises(): +async def test_codec_failures_degrade_to_a_miss_or_a_no_op() -> None: def _boom(_: str) -> str: - raise ValueError("encrypt unavailable") + raise ValueError("codec 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 - - -@pytest.mark.asyncio -async def test_get_is_a_miss_when_decrypt_raises(): - # A blob the decrypt rejects with an exception (e.g. bad ciphertext after key rotation) must read - # as a miss, not propagate the error out of get() and abort fetch(). - def _boom(_: str) -> str | None: - 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" - assert await DualCacheTokenCacheBackend(cache, codec).get("alice", "srv") is None + server = fakeredis.FakeServer() + await _backend(_worker(server), OAuthTokenCacheCodec(encrypt=_boom, decrypt=_boom)).set( + "alice", "srv", OAuthToken(access_token="at"), 60.0 + ) + await fakeredis.FakeAsyncRedis(server=server).set(KEY, json.dumps("whatever")) + assert ( + await _backend(_worker(server), OAuthTokenCacheCodec(encrypt=_boom, decrypt=_boom)).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..a68d454eaca 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 @@ -8,7 +8,7 @@ Unit tests for per-user MCP OAuth token storage: import sys from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Optional +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -22,7 +22,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 +40,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", @@ -294,45 +293,55 @@ class TestMCPPerUserTokenCache: ) @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_per_user_key(self, cache): + import json + + import fakeredis + + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from tests._support.fake_redis import fake_redis_cache + + server = fakeredis.FakeServer() + shared = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=fake_redis_cache(server)) 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", shared), ): - await cache.set("alice", "slack-test", "xoxb-token", ttl=3540) + await cache.set("bob", "github-server", "ghp_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 + raw = fakeredis.FakeAsyncRedis(server=server) + mock_encrypt.assert_called_once_with("ghp_token") + assert await raw.get("mcp:per_user_token:bob:github-server") == json.dumps("encrypted_blob_xyz").encode() + assert await raw.ttl("mcp:per_user_token:bob:github-server") == 3540 @pytest.mark.asyncio - async def test_set_uses_correct_cache_key(self, cache, mock_dual_cache): + async def test_delete_drops_the_token_here_and_in_redis_and_tells_peers(self, cache): + import fakeredis + + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from tests._support.fake_redis import fake_redis_cache + + server = fakeredis.FakeServer() + shared = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=fake_redis_cache(server)) with ( + patch("litellm.proxy.mcp_server.oauth2_token_cache.encrypt_value_helper", return_value="enc"), + patch("litellm.proxy.proxy_server.user_api_key_cache", shared), 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), + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new=AsyncMock(), + ) as publish, ): - 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): + await cache.set("alice", "slack-test", "token", ttl=60) 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 shared.in_memory_cache.get_cache("mcp:per_user_token:alice:slack-test") is None + assert await fakeredis.FakeAsyncRedis(server=server).exists("mcp:per_user_token:alice:slack-test") == 0 + publish.assert_awaited_once_with(cache_key="mcp:per_user_token:alice:slack-test") @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_token.py b/tests/unit/redis/stores/test_mcp_per_user_token.py new file mode 100644 index 00000000000..ad9bb9e8de5 --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_per_user_token.py @@ -0,0 +1,49 @@ +import json +from datetime import timedelta +from typing import Final + +import fakeredis +import pytest + +from litellm.redis.stores.mcp_per_user_token import TransportMcpPerUserTokenStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache + +KEY: Final = "ns:mcp:per_user_token:alice:srv" + + +@pytest.mark.asyncio +async def test_blob_bytes_ttl_and_namespace_match_the_legacy_dual_cache_writer(server: fakeredis.FakeServer) -> None: + store: Final = TransportMcpPerUserTokenStore(fake_redis_cache(server, namespace="ns").transport()) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + + assert await store.put("alice", "srv", 'enc "\u96ea"', ttl=timedelta(seconds=90.9)) is True + assert await raw.get(KEY) == json.dumps('enc "\u96ea"').encode() + assert await raw.ttl(KEY) == 90 + assert await store.get("alice", "srv") == 'enc "\u96ea"' + assert await store.delete("alice", "srv") == 1 + assert await store.get("alice", "srv") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("stored", "expected"), [(b"'legacy-literal'", "legacy-literal"), (b"12345", None)]) +async def test_legacy_literal_and_non_string_values_decode_like_the_dual_cache( + server: fakeredis.FakeServer, stored: bytes, expected: str | None +) -> None: + await fakeredis.FakeAsyncRedis(server=server).set(KEY, stored) + store: Final = TransportMcpPerUserTokenStore(fake_redis_cache(server, namespace="ns").transport()) + assert await store.get("alice", "srv") == expected + + +@pytest.mark.asyncio +async def test_malformed_values_are_decode_failures_and_outages_are_transport_failures( + server: fakeredis.FakeServer, +) -> None: + await fakeredis.FakeAsyncRedis(server=server).set(KEY, b"{not json") + store: Final = TransportMcpPerUserTokenStore(fake_redis_cache(server, namespace="ns").transport()) + malformed: Final = await store.get("alice", "srv") + server.connected = False + outage: Final = await store.get("alice", "srv") + + assert isinstance(malformed, RedisFailure) and malformed.kind == "decode" + assert isinstance(outage, RedisFailure) and outage.kind == "transport"