diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index bc1068d77d0..668e936b8a2 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -3,7 +3,7 @@ import binascii import hashlib import json from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Union, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -39,6 +39,9 @@ from litellm.repositories.verification_token_repository import ( from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPCredentials +if TYPE_CHECKING: + from litellm.types.mcp_server.mcp_server_manager import MCPServer + def _is_global_env_var_scope(scope: Any) -> bool: """``scope="user"`` entries are placeholders the user fills in; everything @@ -1222,6 +1225,91 @@ async def resolve_valid_user_oauth_token( return refreshed +async def resolve_user_oauth_access_token( + user_id: str | None, + server: "MCPServer", + prefetched_creds: dict[str, dict[str, object]] | None = None, +) -> str | None: + """Resolve a user's valid OAuth2 access token for a server: Redis cache, else DB + refresh. + + The egress token-resolution core shared by v1's header builder and the v2 ``OAuthTokenStore`` + adapter. Redis fast-path (skipped when ``prefetched_creds`` is supplied), else a DB read through + ``resolve_valid_user_oauth_token`` (which refreshes an expired token when a ``refresh_token`` is + stored), re-warming the Redis cache with the per-server TTL. Returns ``None`` when there is no + usable token; any error is swallowed to ``None`` so a transient failure reads as "not + authorized" rather than raising. + """ + server_id = getattr(server, "server_id", None) + if not user_id or not server_id: + return None + try: + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + _compute_per_user_token_ttl, + mcp_per_user_token_cache, + ) + + if prefetched_creds is None: + cached_token = await mcp_per_user_token_cache.get(user_id, server_id) + if cached_token is not None: + return cached_token + + prisma_client = None + if prefetched_creds is not None: + cred = prefetched_creds.get(server_id) + else: + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Connect a database to use OAuth2 MCP tools." + ) + cred = await get_user_oauth_credential(prisma_client, user_id, server_id) + + if not cred or not cred.get("access_token"): + return None + + cred = await resolve_valid_user_oauth_token( + user_id=user_id, + server=server, + cred=cred, + prisma_client=prisma_client, + ) + if cred is None: + # Refresh failed or token expired with no usable refresh_token — clear the stale + # Redis entry so the next request doesn't reuse it. + await mcp_per_user_token_cache.delete(user_id, server_id) + return None + + access_token: str = cred["access_token"] + if prefetched_creds is None: + ttl = _compute_per_user_token_ttl( + server, _remaining_token_seconds(cred.get("expires_at")) + ) + await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl) + return access_token + except Exception as e: + verbose_proxy_logger.warning( + "resolve_user_oauth_access_token: failed for user=%s server=%s: %s", + user_id, + server_id, + e, + ) + return None + + +def _remaining_token_seconds(expires_at: str | None) -> int | None: + """Seconds until ``expires_at`` (ISO 8601), or None when absent/past/unparseable.""" + if not expires_at: + return None + try: + exp_dt = datetime.fromisoformat(expires_at) + except (ValueError, TypeError): + return None + if exp_dt.tzinfo is None: + exp_dt = exp_dt.replace(tzinfo=timezone.utc) + remaining = int((exp_dt - datetime.now(timezone.utc)).total_seconds()) + return remaining if remaining > 0 else None + + async def approve_mcp_server( prisma_client: PrismaClient, server_id: str, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v1_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v1_token_store.py index 97d37d7c97d..c5e2f1f266f 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v1_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v1_token_store.py @@ -1,37 +1,49 @@ """v1-backed ``OAuthTokenStore`` source for the ``authorization_code`` mode. -Reads the user's stored access token through v1's ``mcp_per_user_token_cache`` (a Redis-backed, -encrypted-at-rest per-user cache). This is a temporary adapter: step 1b replaces it with a -v2-native token store that also tracks expiry and refresh, behind the same ``OAuthTokenStore`` seam. -It imports v1, so it is kept out of the package ``__init__`` like the rest of the adapter layer. +Resolves the user's access token through v1's egress core (``resolve_user_oauth_access_token``: +Redis cache, else DB read with refresh), so the v2 arm injects exactly the token v1 would, with the +same silent refresh. This is the strangler adapter; step 1b swaps the core for a v2-native store +behind the same ``OAuthTokenStore`` seam. It imports v1, so it is kept out of the package +``__init__`` like the rest of the adapter layer. """ from __future__ import annotations +from collections.abc import Callable +from typing import TYPE_CHECKING from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, ) +if TYPE_CHECKING: + from litellm.types.mcp_server.mcp_server_manager import MCPServer + class V1PerUserTokenStore: - """``OAuthTokenStore`` backed by v1's per-user token cache. + """``OAuthTokenStore`` backed by v1's per-user OAuth egress. - v1 stores only the access token (the cache TTL is its lifetime), so the ``OAuthToken`` carries - no ``expires_at`` or ``refresh_token``: the v2 cache holds it for its default TTL, and the OAuth - challenge drives re-auth once v1's cache drops it. v1's ``get`` swallows errors as a miss, so - this never raises ``TokenStoreUnavailable``. + ``fetch`` looks the server up by id (injected ``server_lookup``) and resolves the token through + v1's ``resolve_user_oauth_access_token``, which refreshes an expired token when a refresh_token + is stored and re-warms the Redis cache. Returns ``None`` when the user has no usable token (the + arm turns that into a challenge); v1's core swallows store errors as a miss, so this never + raises ``TokenStoreUnavailable``. The returned ``OAuthToken`` carries only the access token — + refresh happens inside the core, not via ``RefreshingTokenStore``. """ + def __init__(self, server_lookup: Callable[[str], MCPServer | None]) -> None: + self._server_lookup = server_lookup + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: if not user_id: return None + server = self._server_lookup(server_id) + if server is None: + return None - from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( - mcp_per_user_token_cache, + from litellm.proxy._experimental.mcp_server.db import ( + resolve_user_oauth_access_token, ) - access_token = await mcp_per_user_token_cache.get(user_id, server_id) - if not access_token: - return None - return OAuthToken(access_token=access_token) + access_token = await resolve_user_oauth_access_token(user_id, server) + return OAuthToken(access_token=access_token) if access_token else None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index d7975303802..2ce9054d4b9 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1360,115 +1360,21 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None, ) -> Optional[Dict[str, str]]: - """Look up stored OAuth2 token for (user, server) and return as extra_headers dict. + """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. - Lookup order: - 1. Redis cache (fast path, NaCl-decrypted) — skipped when prefetched_creds supplied - 2. prefetched_creds dict (pre-fetched batch DB query) or fresh DB query - 3. Auto-refresh when the stored token is expired and a refresh_token exists - - Args: - prefetched_creds: Optional dict keyed by server_id with credential payloads. - When provided, the Redis and individual DB lookups are - skipped in favour of the pre-fetched batch result. + Thin wrapper over ``resolve_user_oauth_access_token`` (Redis cache, else DB + refresh); + ``prefetched_creds`` skips the per-server Redis/DB lookups for the batch path. """ - if server.auth_type != MCPAuth.oauth2: + if server.auth_type != MCPAuth.oauth2 or user_api_key_auth is None: return None - if user_api_key_auth is None: - return None - user_id = getattr(user_api_key_auth, "user_id", None) - server_id = getattr(server, "server_id", None) - if not user_id or not server_id: - return None - try: - from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 - get_user_oauth_credential, - resolve_valid_user_oauth_token, - ) - from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415 - _compute_per_user_token_ttl, - mcp_per_user_token_cache, - ) + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + resolve_user_oauth_access_token, + ) - # ── Fast path: Redis cache ──────────────────────────────────────── - # Only used when prefetched_creds is not supplied (individual lookup). - if prefetched_creds is None: - cached_token = await mcp_per_user_token_cache.get(user_id, server_id) - if cached_token is not None: - verbose_logger.debug( - "_get_user_oauth_extra_headers_from_db: Redis hit for user=%s server=%s", - user_id, - server_id, - ) - return {"Authorization": f"Bearer {cached_token}"} - - # ── Slow path: DB lookup ────────────────────────────────────────── - prisma_client = None - if prefetched_creds is not None: - cred = prefetched_creds.get(server_id) - else: - from litellm.proxy.utils import ( # noqa: PLC0415 - get_prisma_client_or_throw, - ) - - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to use OAuth2 MCP tools." - ) - cred = await get_user_oauth_credential( - prisma_client, user_id, server_id - ) - - if not cred or not cred.get("access_token"): - return None - - cred = await resolve_valid_user_oauth_token( - user_id=user_id, - server=server, - cred=cred, - prisma_client=prisma_client, - ) - if cred is None: - # Refresh failed or token expired with no usable refresh_token — - # clear the stale Redis entry so the next request doesn't reuse it. - await mcp_per_user_token_cache.delete(user_id, server_id) - return None - - access_token: str = cred["access_token"] - - # Warm (or re-warm) the Redis cache from the DB result. - # Always write regardless of whether expires_at is present — tokens - # without an expiry are still valid and should be cached using the - # server/default TTL so subsequent requests are fast. - if prefetched_creds is None: - raw_expires = None - expires_at = cred.get("expires_at") - if expires_at: - from datetime import datetime, timezone # noqa: PLC0415 - - try: - exp_dt = datetime.fromisoformat(expires_at) - if exp_dt.tzinfo is None: - exp_dt = exp_dt.replace(tzinfo=timezone.utc) - remaining = int( - (exp_dt - datetime.now(timezone.utc)).total_seconds() - ) - raw_expires = max(remaining, 0) if remaining > 0 else None - except (ValueError, TypeError): - pass - ttl = _compute_per_user_token_ttl(server, raw_expires) - await mcp_per_user_token_cache.set( - user_id, server_id, access_token, ttl - ) - - return {"Authorization": f"Bearer {access_token}"} - except Exception as e: - verbose_logger.warning( - "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for user=%s server=%s: %s", - user_id, - server_id, - e, - ) - return None + token = await resolve_user_oauth_access_token( + getattr(user_api_key_auth, "user_id", None), server, prefetched_creds + ) + return {"Authorization": f"Bearer {token}"} if token else None async def _prefetch_oauth_creds_for_user( user_api_key_auth: Optional[UserAPIKeyAuth], diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v1_token_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v1_token_store.py index 4160e0d0d51..2595fceef09 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v1_token_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v1_token_store.py @@ -6,32 +6,42 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.v1_token_store V1PerUserTokenStore, ) -_GET = ( - "litellm.proxy._experimental.mcp_server.oauth2_token_cache." - "mcp_per_user_token_cache.get" -) +_RESOLVE = "litellm.proxy._experimental.mcp_server.db.resolve_user_oauth_access_token" -async def test_wraps_the_v1_access_token(): - with patch(_GET, new=AsyncMock(return_value="at-123")): - token = await V1PerUserTokenStore().fetch("alice", "s") +def _store_for(server: object) -> V1PerUserTokenStore: + return V1PerUserTokenStore(server_lookup=lambda _server_id: server) + + +async def test_wraps_the_resolved_access_token(): + with patch(_RESOLVE, new=AsyncMock(return_value="at-123")): + token = await _store_for(object()).fetch("alice", "s") assert token is not None and token.access_token == "at-123" async def test_missing_token_is_none(): - with patch(_GET, new=AsyncMock(return_value=None)): - assert await V1PerUserTokenStore().fetch("alice", "s") is None + with patch(_RESOLVE, new=AsyncMock(return_value=None)): + assert await _store_for(object()).fetch("alice", "s") is None -async def test_empty_user_short_circuits_without_hitting_v1(): - get = AsyncMock(return_value="at") - with patch(_GET, new=get): - assert await V1PerUserTokenStore().fetch("", "s") is None - get.assert_not_called() +async def test_empty_user_short_circuits_without_resolving(): + resolve = AsyncMock(return_value="at") + with patch(_RESOLVE, new=resolve): + assert await _store_for(object()).fetch("", "s") is None + resolve.assert_not_called() -async def test_passes_user_and_server_through_to_v1(): - get = AsyncMock(return_value="at") - with patch(_GET, new=get): - await V1PerUserTokenStore().fetch("alice", "srv-1") - get.assert_awaited_once_with("alice", "srv-1") +async def test_unknown_server_is_none_without_resolving(): + resolve = AsyncMock(return_value="at") + store = V1PerUserTokenStore(server_lookup=lambda _server_id: None) + with patch(_RESOLVE, new=resolve): + assert await store.fetch("alice", "missing") is None + resolve.assert_not_called() + + +async def test_passes_user_and_resolved_server_to_the_core(): + server = object() + resolve = AsyncMock(return_value="at") + with patch(_RESOLVE, new=resolve): + await _store_for(server).fetch("alice", "srv-1") + resolve.assert_awaited_once_with("alice", server)