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

1"""Serialize + encrypt boundary for caching an OAuth token in a shared (Redis) cache. 

2 

3Shared cache values contain an encrypted access token and optional identity-binding proof. 

4Refresh tokens remain in the database; cache TTL bounds the access token's lifetime. 

5Legacy bearer-only entries decode without proof and cannot satisfy identity enforcement. 

6""" 

7 

8from __future__ import annotations 

9 

10import json 

11from collections.abc import Callable 

12from dataclasses import dataclass 

13from typing import Final 

14 

15from pydantic import TypeAdapter, ValidationError 

16 

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

18 OAuthToken, 

19) 

20 

21_BOUND_PREFIX: Final = "litellm-bound-oauth-v1:" 

22_BOUND_PAYLOAD: Final = TypeAdapter(dict[str, str]) 

23 

24 

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

26class OAuthTokenCacheCodec: 

27 encrypt: Callable[[str], str] 

28 decrypt: Callable[[str], str | None] 

29 

30 def encode(self, token: OAuthToken) -> str: 

31 if token.identity_binding_proof is not None: 

32 return self.encrypt( 

33 _BOUND_PREFIX 

34 + json.dumps( 

35 { 

36 "access_token": token.access_token, 

37 "identity_binding_proof": token.identity_binding_proof, 

38 } 

39 ) 

40 ) 

41 return self.encrypt(token.access_token) 

42 

43 def decode(self, blob: str) -> OAuthToken | None: 

44 access_token: Final = self.decrypt(blob) 

45 if not access_token: 

46 return None 

47 if access_token.startswith(_BOUND_PREFIX): 

48 try: 

49 payload: Final = _BOUND_PAYLOAD.validate_json(access_token[len(_BOUND_PREFIX) :]) 

50 except ValidationError: 

51 return None 

52 bearer: Final = payload.get("access_token") 

53 proof: Final = payload.get("identity_binding_proof") 

54 if not bearer or not proof: 

55 return None 

56 return OAuthToken(access_token=bearer, identity_binding_proof=proof) 

57 return OAuthToken(access_token=access_token, refresh_token=None)