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

1"""v2-native per-user OAuth token read store for the ``authorization_code`` mode. 

2 

3The raw "inner" store that ``RefreshingTokenStore`` and ``CachedOAuthTokenStore`` wrap: it reads the 

4user's persisted credential and returns a typed ``OAuthToken`` (access token, epoch expiry, refresh 

5token), validating the decoded credential blob at this boundary so no ``Any`` leaks past it. It does 

6not cache or refresh - those are the decorators. This replaces ``V1PerUserTokenStore`` (which handed 

7the whole read + cache + refresh to v1's core) as step 1b: the ``read_credential`` collaborator is 

8injected, so the DB/decoding plumbing stays testable and out of this seam. 

9""" 

10 

11from __future__ import annotations 

12 

13from collections.abc import Awaitable, Callable, Mapping 

14from datetime import datetime, timezone 

15from typing import Final 

16 

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

18 OAuthToken, 

19) 

20 

21CredentialReader = Callable[[str, str], Awaitable["Mapping[str, object] | None"]] 

22 

23 

24def _iso_to_epoch(expires_at: str) -> float | None: 

25 try: 

26 dt = datetime.fromisoformat(expires_at) 

27 except ValueError: 

28 return None 

29 # A timezone-naive expiry is stored as UTC (db.py writes ``datetime.now(timezone.utc)``), 

30 # so anchor it to UTC before ``.timestamp()`` - otherwise a non-UTC host would read it as 

31 # local time and skew the expiry, diverging from v1's ``_remaining_token_seconds``. 

32 if dt.tzinfo is None: 

33 dt = dt.replace(tzinfo=timezone.utc) 

34 return dt.timestamp() 

35 

36 

37def _to_scopes(raw: object) -> tuple[str, ...]: 

38 if isinstance(raw, (list, tuple)): 

39 return tuple(s for s in raw if isinstance(s, str)) 

40 return () 

41 

42 

43def _to_oauth_token(payload: Mapping[str, object]) -> OAuthToken | None: 

44 access_token: Final = payload.get("access_token") 

45 if not isinstance(access_token, str): 

46 return None 

47 refresh_token: Final = payload.get("refresh_token") 

48 expires_at: Final = payload.get("expires_at") 

49 binding_proof: Final = payload.get("identity_binding_proof") 

50 return OAuthToken( 

51 access_token=access_token, 

52 expires_at=_iso_to_epoch(expires_at) if isinstance(expires_at, str) else None, 

53 refresh_token=refresh_token if isinstance(refresh_token, str) else None, 

54 scopes=_to_scopes(payload.get("scopes")), 

55 identity_binding_proof=binding_proof if isinstance(binding_proof, str) else None, 

56 ) 

57 

58 

59class V2PerUserTokenStore: 

60 """``OAuthTokenStore`` that reads the user's persisted authorization_code credential, typed. 

61 

62 The injected ``read_credential`` returns the decoded credential payload for a ``(user, server)`` 

63 pair, or ``None`` when the user has not completed OAuth. A backing-store outage surfaces as 

64 ``TokenStoreUnavailable`` from the reader, which the arm turns into a challenge rather than a 

65 500, so ``fetch`` lets it propagate. Refresh is the wrapping ``RefreshingTokenStore``'s job, so 

66 the returned token carries ``expires_at`` and ``refresh_token`` for it to act on. 

67 """ 

68 

69 def __init__(self, read_credential: CredentialReader) -> None: 

70 self._read_credential = read_credential 

71 

72 async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: 

73 if not user_id: 

74 return None 

75 payload: Final = await self._read_credential(user_id, server_id) 

76 if payload is None: 

77 return None 

78 return _to_oauth_token(payload)