Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/byok_credential_cache.py: 100%

1"""Per-worker cache of stored BYOK credentials, keyed so peer workers can evict it over the auth cache pub/sub.""" 

2 

3from dataclasses import dataclass 

4from typing import Final 

5 

6from litellm.caching.in_memory_cache import InMemoryCache 

7from litellm.constants import MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE, MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS 

8 

9_CACHE_KEY_PREFIX: Final = "mcp_byok_credential" 

10 

11 

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

13class CachedByokCredential: 

14 credential: str | None 

15 

16 

17byok_credential_cache: Final = InMemoryCache( 

18 max_size_in_memory=MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE, 

19 default_ttl=MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS, 

20) 

21 

22 

23def byok_credential_cache_key(user_id: str, server_id: str) -> str: 

24 return f"{_CACHE_KEY_PREFIX}:{user_id}:{server_id}" 

25 

26 

27def get_cached_byok_credential(user_id: str, server_id: str) -> CachedByokCredential | None: 

28 cached: Final = byok_credential_cache.get_cache( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # InMemoryCache is untyped 

29 byok_credential_cache_key(user_id, server_id) 

30 ) 

31 return cached if isinstance(cached, CachedByokCredential) else None 

32 

33 

34def cache_byok_credential(user_id: str, server_id: str, credential: str | None) -> None: 

35 byok_credential_cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped 

36 byok_credential_cache_key(user_id, server_id), 

37 CachedByokCredential(credential=credential), 

38 )