feat(mcp): migrate authorization_code MCP to the v2 resolver (single-replica) [1/2]

litellm/proxy/_experimental/mcp_server/db.py

@@ -3,7 +3,7 @@
 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.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,

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -63,9 +63,16 @@
 )
 from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
     raise_public,
+    raise_user_oauth_challenge,
     to_server_spec,
     to_subject,
 )
+from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import (
+    LazyPerUserOAuthTokenStore,
+)
+from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
+    AuthorizationCodeConfig,
+)
 from litellm.proxy._experimental.mcp_server.utils import (
     MCP_TOOL_PREFIX_SEPARATOR,
     MCPMissingUserEnvVarsError,
@@ -186,6 +193,10 @@ def _should_strip_caller_authorization(
     Strip rules:
     - **M2M (client_credentials) servers**: never forward the caller's
       ``Authorization`` — the proxy fetches its own upstream token.
+    - **Migrated per-user OAuth (authorization_code) servers**: never forward
+      the caller's ``Authorization`` — the v2 resolver injects the stored
+      per-user token, so a caller-supplied bearer cannot override another
+      user's stored credential. Delegate / pass-through keep forwarding it.
     - **OAuth pass-through servers**: strip when the ``Authorization``
       header is actually the LiteLLM API key — either because admission
       validated it (``user_api_key_auth.api_key`` is set) and the caller
@@ -198,6 +209,15 @@ def _should_strip_caller_authorization(
     """
     if mcp_server.has_client_credentials:
         return True
+    if (
+        mcp_server.auth_type == MCPAuth.oauth2
+        and to_server_spec(mcp_server) is not None
+    ):
+        # Migrated per-user OAuth (authorization_code): the v2 resolver injects the
+        # stored token, so a caller-forwarded Authorization must not be forwarded
+        # upstream — it would override another user's stored credential. Delegate and
+        # pass-through return None from to_server_spec and keep forwarding the bearer.
+        return True
     if not mcp_server.is_oauth_passthrough:
         return False
 
@@ -217,6 +237,18 @@ def _should_strip_caller_authorization(
     )
 
 
+def _without_authorization(
+    headers: Optional[dict[str, str]],
+) -> Optional[dict[str, str]]:
+    """A copy of ``headers`` with any ``Authorization`` key removed (case-insensitive), or
+    None if nothing remains. Drops only the credential, keeping other forwarded headers.
+    """
+    if not headers:
+        return None
+    filtered = {k: v for k, v in headers.items() if k.lower() != "authorization"}
+    return filtered or None
+
+
 def _extract_upstream_auth_failure(
     exc: BaseException,
 ) -> Optional[Tuple[int, Optional[str]]]:
@@ -523,7 +555,9 @@ def _resolve_oauth2_flow(
         return None
 
     def __init__(self, cred_provider: Optional[UpstreamCredentialProvider] = None):
-        self._cred_provider = cred_provider or UpstreamCredentialProvider()
+        self._cred_provider = cred_provider or UpstreamCredentialProvider(
+            oauth_token_store=LazyPerUserOAuthTokenStore(self.get_mcp_server_by_id)
+        )
         self.registry: Dict[str, MCPServer] = {}
         self.config_mcp_servers: Dict[str, MCPServer] = {}
         """
@@ -1934,6 +1968,7 @@ async def _create_mcp_client(
         stdio_env: Optional[Dict[str, str]] = None,
         subject_token: Optional[str] = None,
         user_api_key_auth: Optional[UserAPIKeyAuth] = None,
+        cred_provider: Optional[UpstreamCredentialProvider] = None,
     ) -> MCPClient:
         """
         Create an MCPClient instance for the given server.
@@ -1957,11 +1992,17 @@ async def _create_mcp_client(
         """
         transport = server.transport or MCPTransport.sse
         spec = None if transport == MCPTransport.stdio else to_server_spec(server)
-        # A per-request override is the caller-supplied credential v1 turns into the upstream
-        # auth, so it must win; defer those to v1 (this defer falls away once the per-user modes
-        # stop writing mcp_auth_header). An inbound header already in extra_headers is handled on
-        # the v2 path below, not here.
-        if spec is not None and mcp_auth_header:
+        provider = cred_provider or self._cred_provider
+        # A caller-supplied per-request override (mcp_auth_header / x-mcp-*) defers to the v1 path
+        # so it wins - except for authorization_code, whose per-user token the v2 resolver owns. A
+        # caller must not be able to substitute another user's stored credential, so we keep the v2
+        # spec and ignore the override there; the REST tools preview supplies its not-yet-persisted
+        # token through the resolver (cred_provider), never this path.
+        if (
+            spec is not None
+            and mcp_auth_header
+            and not isinstance(spec.config, AuthorizationCodeConfig)
+        ):
             spec = None
         auth_value = (
             await resolve_mcp_auth(server, mcp_auth_header, subject_token=subject_token)
@@ -2039,7 +2080,7 @@ async def _create_mcp_client(
             server_url = server.url or ""
 
             if spec is not None:
-                match await self._cred_provider.resolve_credentials(
+                match await provider.resolve_credentials(
                     to_subject(user_api_key_auth, subject_token), spec
                 ):
                     case Ok(auth):
@@ -2059,6 +2100,10 @@ async def _create_mcp_client(
                         ):
                             resolved_auth = None
                     case Error(err):
+                        if err.tag == "unauthorized":
+                            # The arm signals a missing per-user token semantically; raise the
+                            # per-server OAuth challenge here, where the full MCPServer is in hand.
+                            raise_user_oauth_challenge(server)
                         raise_public(err)
                 return MCPClient(
                     server_url=server_url,
@@ -3499,6 +3544,16 @@ async def _call_regular_mcp_tool(
                 extra_headers = None
             else:
                 extra_headers = oauth2_headers
+                # Migrated authorization_code: the v2 resolver injects the stored per-user
+                # token, so drop the caller-forwarded Authorization (apply-if-absent would
+                # otherwise let it shadow the resolved token). Delegate keeps it. Centralized
+                # via _should_strip_caller_authorization to match _prepare_mcp_server_headers.
+                if extra_headers and _should_strip_caller_authorization(
+                    mcp_server=mcp_server,
+                    raw_headers=raw_headers,
+                    user_api_key_auth=user_api_key_auth,
+                ):
+                    extra_headers = _without_authorization(extra_headers)
 
         if mcp_server.extra_headers and raw_headers:
             if extra_headers is None:
@@ -3670,6 +3725,22 @@ def _candidate_matches_server_name(candidate: MCPServer) -> bool:
 
         return mcp_server
 
+    async def has_user_oauth_token(
+        self, server: MCPServer, user_api_key_auth: Optional[UserAPIKeyAuth]
+    ) -> bool:
+        """Whether the v2 resolver can produce a per-user token for this server right now.
+
+        This is the preemptive 401's existence check, routed through the same resolver that drives
+        the egress so every authorization_code resolution (egress and the discovery challenge) runs
+        through v2. Returns False for a server the resolver does not own (a None spec).
+        """
+        spec = to_server_spec(server)
+        if spec is None:
+            return False
+        return await self._cred_provider.has_user_token(
+            to_subject(user_api_key_auth, None), spec
+        )
+
     async def _resolve_oauth2_headers_for_tool_call(
         self,
         mcp_server: MCPServer,
@@ -3684,6 +3755,12 @@ async def _resolve_oauth2_headers_for_tool_call(
         ):
             return oauth2_headers
 
+        if to_server_spec(mcp_server) is not None:
+            # Migrated to v2: the resolver owns this server's per-user token (inject or fail-closed
+            # 401). Building it into extra_headers here would let the v2 graft defer to it and
+            # shadow the resolver, double-resolving and hiding the per-server challenge.
+            return oauth2_headers
+
         user_id = getattr(user_api_key_auth, "user_id", None)
         if not user_id:
             return oauth2_headers

litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py

@@ -20,6 +20,7 @@
 
 from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
     ApiKeyConfig,
+    AuthorizationCodeConfig,
     CredError,
     NoneConfig,
     ServerSpec,
@@ -61,8 +62,9 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]:
     Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with
     an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is
     explicitly mapped or explicitly deferred, rather than silently falling through to v1. Live
-    modes: ``none`` and the static-header family (``api_key`` plus the Authorization schemes),
-    all shared-key; every other mode returns None and stays on v1.
+    modes: ``none``, the static-header family (``api_key`` plus the Authorization schemes,
+    all shared-key), and ``oauth2`` per-user tokens (``authorization_code``); client_credentials
+    (M2M), delegated/passthrough oauth2, token exchange, and SigV4 return None and stay on v1.
     """
     if server.is_byok:
         return (
@@ -89,8 +91,17 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]:
             return _shared_key_spec(
                 server, resource, "Authorization", "Basic", encode=True
             )
-        case MCPAuth.oauth2 | MCPAuth.oauth2_token_exchange | MCPAuth.aws_sigv4:
-            return None  # OAuth grants and SigV4 are not migrated yet -> defer to v1
+        case MCPAuth.oauth2:
+            if server.needs_user_oauth_token and not server.delegate_auth_to_upstream:
+                return ServerSpec(
+                    server_id=server.server_id,
+                    resource=resource,
+                    config=AuthorizationCodeConfig(),
+                )
+            # client_credentials (M2M) and delegate/passthrough oauth2 stay on v1
+            return None
+        case MCPAuth.oauth2_token_exchange | MCPAuth.aws_sigv4:
+            return None  # token exchange and SigV4 are not migrated yet -> defer to v1
     assert_never(auth_type)
 
 
@@ -147,3 +158,25 @@ def raise_public(error: CredError) -> NoReturn:
         case "not_implemented":
             raise HTTPException(status_code=501, detail=error.summary)
     assert_never(error.tag)
+
+
+def raise_user_oauth_challenge(server: MCPServer) -> NoReturn:
+    """Raise the 401 an ``authorization_code`` server returns at egress when the user has no token.
+
+    Points at the server's RFC 9728 Protected Resource Metadata (``resource_metadata``), which names
+    the upstream authorization server the client must complete OAuth with. The URL is per-server and
+    relative, so it resolves against the caller's own host (correct even behind a reverse proxy)
+    without needing request context. The listing-phase 401 still emits the RFC 8414 ``authorization_uri``
+    form pending the format unification; both target the same server, so the difference is cosmetic.
+    """
+    from litellm.proxy.utils import get_server_root_path  # noqa: PLC0415
+
+    root = get_server_root_path()
+    prefix = "" if root == "/" else root
+    name = server.alias or server.server_name or server.name or server.server_id
+    resource_metadata = f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}"
+    raise HTTPException(
+        status_code=401,
+        detail="Unauthorized",
+        headers={"WWW-Authenticate": f'Bearer resource_metadata="{resource_metadata}"'},
+    )