Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/outbound_credentials/dual_cache_token_backend.py: 51%

1"""Cross-replica ``TokenCacheBackend``: stores the token in LiteLLM's shared ``DualCache``. 

2 

3Plugs into the foundation's ``CachedOAuthTokenStore`` via the ``TokenCacheBackend`` seam. The token is 

4encrypted + serialized by the injected codec and written under a per-``(user, server)`` key with the 

5given TTL, so every worker reads one refresh rather than each re-reading and re-refreshing - matching 

6v1's ``MCPPerUserTokenCache`` (same NaCl encryption and key, so a token cached by either is readable by 

7the other across the cutover). A missing or undecryptable entry reads as a miss. 

8""" 

9 

10from __future__ import annotations 

11 

12from dataclasses import KW_ONLY, dataclass 

13from typing import Final, Protocol 

14 

15from litellm._logging import verbose_logger 

16from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( 

17 OAuthToken, 

18) 

19from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import ( 

20 OAuthTokenCacheCodec, 

21) 

22 

23 

24class AsyncCache(Protocol): 

25 """The slice of LiteLLM's ``DualCache`` this backend needs (Redis-backed, shared across workers).""" 

26 

27 async def async_get_cache(self, key: str) -> object | None: ... 27 ↛ exitline 27 didn't return from function 'async_get_cache' because

28 

29 async def async_set_cache(self, key: str, value: str, ttl: float | None = None) -> None: ... 29 ↛ exitline 29 didn't return from function 'async_set_cache' because

30 

31 async def async_delete_cache(self, key: str) -> None: ... 31 ↛ exitline 31 didn't return from function 'async_delete_cache' because

32 

33 

34@dataclass(frozen=True, slots=True) 

35class DualCacheTokenCacheBackend: 

36 """Every method degrades a cache or codec failure to its safe value - ``get`` to a miss 

37 (``None``), ``set``/``delete`` to a no-op - so a Redis outage or an undecryptable entry reads as a 

38 cache miss rather than a request error, matching v1 and this layer's "boundary failure = miss" 

39 contract. The guarantee holds here regardless of whether the injected cache/codec also swallow. 

40 """ 

41 

42 cache: AsyncCache 

43 codec: OAuthTokenCacheCodec 

44 _: KW_ONLY 

45 key_prefix: str = "mcp:per_user_token:" 

46 

47 def _key(self, user_id: str, server_id: str) -> str: 

48 return f"{self.key_prefix}{user_id}:{server_id}" 

49 

50 async def get(self, user_id: str, server_id: str) -> OAuthToken | None: 

51 try: 

52 blob: Final = await self.cache.async_get_cache(self._key(user_id, server_id)) 

53 return self.codec.decode(blob) if isinstance(blob, str) else None 

54 except Exception as exc: # noqa: BLE001 

55 verbose_logger.debug("MCP per-user token cache get failed (miss): %s", exc) 

56 return None 

57 

58 async def set(self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float) -> None: 

59 if ttl_seconds <= 0: 

60 return 

61 try: 

62 await self.cache.async_set_cache( 

63 self._key(user_id, server_id), 

64 self.codec.encode(token), 

65 ttl=ttl_seconds, 

66 ) 

67 except Exception as exc: # noqa: BLE001 

68 verbose_logger.debug("MCP per-user token cache set failed (ignored): %s", exc) 

69 

70 async def delete(self, user_id: str, server_id: str) -> None: 

71 try: 

72 await self.cache.async_delete_cache(self._key(user_id, server_id)) 

73 except Exception as exc: # noqa: BLE001 

74 verbose_logger.debug("MCP per-user token cache delete failed (ignored): %s", exc)