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/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b9e0379e445..c34b86170fa 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -63,9 +63,16 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( ) 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 @@ class MCPServerManager: 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 @@ class MCPServerManager: 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 @@ class MCPServerManager: """ 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 @@ class MCPServerManager: 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 @@ class MCPServerManager: ): 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 @@ class MCPServerManager: 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 @@ class MCPServerManager: 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 @@ class MCPServerManager: ): 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 diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 91876df762f..b3e1cd844d8 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -20,6 +20,7 @@ from typing_extensions import assert_never 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}"'}, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py new file mode 100644 index 00000000000..04b7a54aaa0 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py @@ -0,0 +1,121 @@ +"""v2-native refresher for the ``authorization_code`` mode: the refresh_token grant, then persist. + +Mints a fresh access token from a stored refresh_token by POSTing the RFC 6749 refresh_token grant to +the server's token endpoint, persists the rotated triple, and returns the new typed ``OAuthToken`` for +``RefreshingTokenStore`` to cache. The HTTP post and the persist are injected, so the orchestration +and the (untyped) response parsing stay testable without a live IdP or DB. Replaces v1's +``refresh_user_oauth_token`` as part of step 1b; rotation safety - one refresh per (user, server) +across replicas - is the wrapping store's distributed single-flight, not this refresher's concern. +""" + +from __future__ import annotations + +import time +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Protocol + +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 + +ServerLookup = Callable[[str], "MCPServer | None"] +TokenEndpointPost = Callable[ + [str, dict[str, str]], Awaitable["dict[str, object] | None"] +] + + +class CredentialPersist(Protocol): + async def __call__( + self, + user_id: str, + server_id: str, + access_token: str, + refresh_token: str | None, + expires_in: int | None, + scopes: tuple[str, ...] | None, + ) -> None: ... + + +def _parse_expires_in(raw: object) -> int | None: + if isinstance(raw, bool): + return None + if isinstance(raw, int): + return raw + if isinstance(raw, str): + try: + return int(raw) + except ValueError: + return None + return None + + +def _parse_scopes(raw: object) -> tuple[str, ...] | None: + return tuple(raw.split()) if isinstance(raw, str) and raw else None + + +class AuthorizationCodeRefresher: + """``TokenRefresher`` for authorization_code: refresh_token grant against the server, then persist. + + ``token_endpoint`` POSTs the OAuth form and returns the parsed JSON body (``None`` on any + transport/HTTP failure, mirroring v1: a failed refresh is a miss, not a 500). ``persist`` writes + the rotated triple for ``(user, server)`` - the v1 ``store_user_oauth_credential`` write, which + stays. Returns ``None`` (the arm challenges) when there is no refresh_token, the server lacks a + token endpoint, or the grant fails; never a stale or partial token. A rotated refresh_token from + the response replaces the old one; an omitted one is carried forward, as are the recorded scopes + when the response omits ``scope``. + """ + + def __init__( + self, + server_lookup: ServerLookup, + token_endpoint: TokenEndpointPost, + persist: CredentialPersist, + *, + clock: Callable[[], float] = time.time, + ) -> None: + self._server_lookup = server_lookup + self._token_endpoint = token_endpoint + self._persist = persist + self._clock = clock + + async def refresh( + self, user_id: str, server_id: str, token: OAuthToken + ) -> OAuthToken | None: + if token.refresh_token is None: + return None + server = self._server_lookup(server_id) + if server is None or not server.token_url: + return None + + form = { + "grant_type": "refresh_token", + "refresh_token": token.refresh_token, + **({"client_id": server.client_id} if server.client_id else {}), + **({"client_secret": server.client_secret} if server.client_secret else {}), + } + body = await self._token_endpoint(server.token_url, form) + if body is None: + return None + access_token = body.get("access_token") + if not isinstance(access_token, str) or not access_token: + return None + + rotated = body.get("refresh_token") + new_refresh = ( + rotated if isinstance(rotated, str) and rotated else token.refresh_token + ) + expires_in = _parse_expires_in(body.get("expires_in")) + scopes = _parse_scopes(body.get("scope")) or token.scopes + + await self._persist( + user_id, server_id, access_token, new_refresh, expires_in, scopes or None + ) + return OAuthToken( + access_token=access_token, + expires_at=self._clock() + expires_in if expires_in is not None else None, + refresh_token=new_refresh, + scopes=scopes, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py index 1473d15da5e..1c089f2a931 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py @@ -31,15 +31,20 @@ class OAuthToken: with each mode); it is never minted into a header directly. ``repr`` masks both secrets so a stray log line cannot leak them (the values are still plain ``str`` for the header path, since ``SecretStr`` resolves as unknown under this repo's basedpyright). + + ``scopes`` is the recorded grant. A refresh response that omits ``scope`` (RFC 6749 §5.1: an + omitted ``scope`` means unchanged) carries the prior value forward, so a refresh never silently + drops it; the resolver itself does not read it. """ access_token: str expires_at: float | None = None refresh_token: str | None = None + scopes: tuple[str, ...] = () def __repr__(self) -> str: has_refresh = self.refresh_token is not None - return f"OAuthToken(access_token=***, expires_at={self.expires_at!r}, has_refresh_token={has_refresh})" + return f"OAuthToken(access_token=***, expires_at={self.expires_at!r}, has_refresh_token={has_refresh}, scopes={self.scopes!r})" class TokenStoreUnavailable(Exception): diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py new file mode 100644 index 00000000000..7ffde55c2c1 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -0,0 +1,129 @@ +"""Composition root for the v2-native authorization_code per-user OAuth token store (step 1b). + +Assembles ``Cached(Refreshing(V2PerUserTokenStore))`` and replaces ``V1PerUserTokenStore`` in the +resolver. The runtime collaborators (DB, HTTP) are LiteLLM globals not ready at import time, so the +chain is built lazily on first use. The cache and refresh coordinator use the foundation's in-process +defaults (correct for a single replica); the cross-replica path is layered on separately. The DB +read/refresh-grant/persist collaborators acquire their globals per call, mirroring v1's lazy-import +pattern. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import ( + AuthorizationCodeRefresher, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + CachedOAuthTokenStore, + OAuthToken, + RefreshingTokenStore, + TokenStoreUnavailable, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import ( + V2PerUserTokenStore, +) + +if TYPE_CHECKING: + from litellm.types.mcp_server.mcp_server_manager import MCPServer + +# A token with no declared expiry is cached for this long; one with an expiry is cached until then. +_DEFAULT_TTL_SECONDS = 300.0 + +ServerLookup = Callable[[str], "MCPServer | None"] + + +async def _read_credential(user_id: str, server_id: str) -> dict[str, object] | None: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + get_user_oauth_credential, + ) + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 + + if prisma_client is None: + raise TokenStoreUnavailable("Database not connected") + return await get_user_oauth_credential(prisma_client, user_id, server_id) + + +async def _persist_credential( + user_id: str, + server_id: str, + access_token: str, + refresh_token: str | None, + expires_in: int | None, + scopes: tuple[str, ...] | None, +) -> None: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + store_user_oauth_credential, + ) + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 + + if prisma_client is None: + return + await store_user_oauth_credential( + prisma_client=prisma_client, + user_id=user_id, + server_id=server_id, + access_token=access_token, + refresh_token=refresh_token, + expires_in=expires_in, + scopes=list(scopes) if scopes else None, + skip_byok_guard=True, + ) + + +async def _post_token_endpoint( + url: str, form: dict[str, str] +) -> dict[str, object] | None: + from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 + get_async_httpx_client, # pyright: ignore + ) + from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 + + # litellm's httpx handler and httpx.Response are only partially typed; the IdP returns a JSON + # object and the refresher validates each field, so the untyped boundary is contained here. + provider = httpxSpecialProvider.Oauth2Check + headers = {"Accept": "application/json"} + # A failed refresh is a miss, not a 500 (matches v1), so any error becomes None. + try: + client = get_async_httpx_client(llm_provider=provider) # pyright: ignore + response = await client.post(url, headers=headers, data=form) # pyright: ignore + response.raise_for_status() # pyright: ignore + body: dict[str, object] = response.json() # pyright: ignore + except Exception as exc: # noqa: BLE001 + verbose_logger.warning("MCP OAuth refresh request failed: %s", exc) + return None + else: + return body # pyright: ignore + + +def build_per_user_oauth_token_store( + server_lookup: ServerLookup, +) -> CachedOAuthTokenStore: + refresher = AuthorizationCodeRefresher( + server_lookup, _post_token_endpoint, _persist_credential + ) + # Cache and refresh coordinator use the foundation's in-process defaults (a single replica needs + # no shared cache or lock); the cross-replica path is layered on separately. + refreshing = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher) + return CachedOAuthTokenStore(refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS) + + +class LazyPerUserOAuthTokenStore: + """``OAuthTokenStore`` that builds the v2-native chain on first ``fetch``. + + The chain's cache/lock collaborators are LiteLLM runtime globals not available when the resolver + is constructed at import time, so construction is deferred to the first request (by when they are + wired). Built once, then reused. + """ + + def __init__(self, server_lookup: ServerLookup) -> None: + self._server_lookup = server_lookup + self._store: CachedOAuthTokenStore | None = None + + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + if self._store is None: + self._store = build_per_user_oauth_token_store(self._server_lookup) + return await self._store.fetch(user_id, server_id) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/presented_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/presented_token_store.py new file mode 100644 index 00000000000..c88d31dd6bc --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/presented_token_store.py @@ -0,0 +1,26 @@ +"""One-shot ``OAuthTokenStore`` for the create/test tools preview. + +The preview tests an unsaved server, so no per-user credential is persisted yet. The operator holds +the just-authorized token; this serves it through the same v2 resolver path runtime uses for the +stored token, so the preview never relies on the caller-credential-override path that +``_create_mcp_client`` refuses for ``authorization_code``. It backs a single preview call, so it +returns its one token regardless of the lookup key. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, +) + + +@dataclass(frozen=True, slots=True) +class PresentedOAuthTokenStore: + """Serves one in-hand token for the single preview call it backs (no DB, no cache).""" + + token: OAuthToken + + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + return self.token diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 969bbf01ec8..87fa66aeab9 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -7,9 +7,9 @@ no precedence cascade. It is wildcard-free with an `assert_never` tail, so addin an arm fails the type gate (basedpyright `reportMatchNotExhaustive`); a bypassed gate fails loudly at runtime instead of returning `None`. -`none` and `api_key` (shared-key source) are live; the remaining arms are `not_implemented` -stubs that each land in a follow-up PR with their injected seam. The self-contained arms read -straight from the config and need no collaborator. Pure v2: no imports from v1. +`none` and `api_key` (shared-key source) are live, as is `authorization_code`, which reads the +user's token from the injected `OAuthTokenStore`. The remaining arms are `not_implemented` stubs +that each land in a follow-up PR with their seam. Pure v2: no imports from v1. """ from __future__ import annotations @@ -21,6 +21,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth impo NoOpAuth, StaticHeaderAuth, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, + OAuthTokenStore, + TokenStoreUnavailable, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Error, Ok, @@ -43,13 +48,26 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ) +class _NullOAuthTokenStore: + """Fail-closed default: with no token store wired, every user reads as not authorized.""" + + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + return None + + class UpstreamCredentialProvider: """Produces the one `httpx.Auth` for a `(subject, upstream)` pair, per declared mode. - Collaborators (the per-mode credential stores and token fetchers) are injected as each arm - is built; the live `none` and `api_key`-shared arms read from the config and need none. + Collaborators (the per-mode credential stores and token fetchers) are injected as each arm is + built; the live `none` and `api_key`-shared arms read from the config and need none, while + `authorization_code` reads the user's token from the injected `OAuthTokenStore`. """ + def __init__(self, oauth_token_store: OAuthTokenStore | None = None) -> None: + self._oauth_token_store: OAuthTokenStore = ( + oauth_token_store or _NullOAuthTokenStore() + ) + async def resolve_credentials( self, subject: Subject, server: ServerSpec ) -> Result[httpx.Auth, CredError]: @@ -65,11 +83,21 @@ class UpstreamCredentialProvider: case TokenExchangeConfig(): return _not_implemented(AuthSpecKind.token_exchange) case AuthorizationCodeConfig(): - return _not_implemented(AuthSpecKind.authorization_code) + return await self._authorization_code(subject, server) case AwsSigV4Config(): return _not_implemented(AuthSpecKind.aws_sigv4) assert_never(server.config) + async def has_user_token(self, subject: Subject, server: ServerSpec) -> bool: + """Whether a usable per-user token exists for this server (the preemptive 401's check). + + Reads from the same per-user store as the ``authorization_code`` arm, so the discovery + challenge and the egress agree on whether the user is authorized. Returns a typed ``bool`` + (no ``httpx.Auth``), unlike ``resolve_credentials``. A non-per-user mode has no token in the + store, so it reads as False without a per-mode branch here. + """ + return await self._authz_token(subject, server) is not None + def _api_key(self, config: ApiKeyConfig) -> Result[httpx.Auth, CredError]: match config.key_source: case SharedKey() as source: @@ -86,6 +114,37 @@ class UpstreamCredentialProvider: ) assert_never(config.key_source) + async def _authorization_code( + self, subject: Subject, server: ServerSpec + ) -> Result[StaticHeaderAuth, CredError]: + token = await self._authz_token(subject, server) + if token is None: + return Error( + CredError.of_unauthorized( + "Authorization required: complete the OAuth flow for this server." + ) + ) + return Ok( + StaticHeaderAuth( + f"Bearer {token.access_token}", header_name="Authorization" + ) + ) + + async def _authz_token( + self, subject: Subject, server: ServerSpec + ) -> OAuthToken | None: + """The user's authorization_code token, or None when absent or the store is unreachable. + + A store outage is mapped to None (the OAuth challenge), not raised, so a transient outage + does not 500; it is the store, not this resolver, that declines to cache the failure. + """ + try: + return await self._oauth_token_store.fetch( + subject.subject_id, server.server_id + ) + except TokenStoreUnavailable: + return None + def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: return Error( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py new file mode 100644 index 00000000000..f1b68042c94 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py @@ -0,0 +1,75 @@ +"""v2-native per-user OAuth token read store for the ``authorization_code`` mode. + +The raw "inner" store that ``RefreshingTokenStore`` and ``CachedOAuthTokenStore`` wrap: it reads the +user's persisted credential and returns a typed ``OAuthToken`` (access token, epoch expiry, refresh +token), validating the decoded credential blob at this boundary so no ``Any`` leaks past it. It does +not cache or refresh - those are the decorators. This replaces ``V1PerUserTokenStore`` (which handed +the whole read + cache + refresh to v1's core) as step 1b: the ``read_credential`` collaborator is +injected, so the DB/decoding plumbing stays testable and out of this seam. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from datetime import datetime, timezone + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, +) + +CredentialReader = Callable[[str, str], Awaitable["dict[str, object] | None"]] + + +def _iso_to_epoch(expires_at: str) -> float | None: + try: + dt = datetime.fromisoformat(expires_at) + except ValueError: + return None + # A timezone-naive expiry is stored as UTC (db.py writes ``datetime.now(timezone.utc)``), + # so anchor it to UTC before ``.timestamp()`` - otherwise a non-UTC host would read it as + # local time and skew the expiry, diverging from v1's ``_remaining_token_seconds``. + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt.timestamp() + + +def _to_scopes(raw: object) -> tuple[str, ...]: + if isinstance(raw, (list, tuple)): + return tuple(s for s in raw if isinstance(s, str)) + return () + + +def _to_oauth_token(payload: dict[str, object]) -> OAuthToken | None: + access_token = payload.get("access_token") + if not isinstance(access_token, str): + return None + refresh_token = payload.get("refresh_token") + expires_at = payload.get("expires_at") + return OAuthToken( + access_token=access_token, + expires_at=_iso_to_epoch(expires_at) if isinstance(expires_at, str) else None, + refresh_token=refresh_token if isinstance(refresh_token, str) else None, + scopes=_to_scopes(payload.get("scopes")), + ) + + +class V2PerUserTokenStore: + """``OAuthTokenStore`` that reads the user's persisted authorization_code credential, typed. + + The injected ``read_credential`` returns the decoded credential payload for a ``(user, server)`` + pair, or ``None`` when the user has not completed OAuth. A backing-store outage surfaces as + ``TokenStoreUnavailable`` from the reader, which the arm turns into a challenge rather than a + 500, so ``fetch`` lets it propagate. Refresh is the wrapping ``RefreshingTokenStore``'s job, so + the returned token carries ``expires_at`` and ``refresh_token`` for it to act on. + """ + + def __init__(self, read_credential: CredentialReader) -> None: + self._read_credential = read_credential + + async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + if not user_id: + return None + payload = await self._read_credential(user_id, server_id) + if payload is None: + return None + return _to_oauth_token(payload) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 7b4f1e13a52..0d71174e9d7 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1034,8 +1034,54 @@ if MCP_AVAILABLE: None if server_model.has_client_credentials else oauth2_headers ) + # Interactive authorization_code tools preview: the operator holds a just-authorized + # token but it is not persisted yet. Resolve it through the v2 resolver via a one-shot + # presented store - the same path runtime uses for the stored token - rather than the + # caller-override path _create_mcp_client refuses for authorization_code. The bare token + # becomes the upstream credential, so it is not also forwarded as a caller header. Gated + # to the v2-mapped oauth2 case (to_server_spec non-None); M2M (client_credentials), + # delegate/passthrough, and token-exchange are unaffected. + from litellm.proxy._experimental.mcp_server.outbound_credentials import ( # noqa: PLC0415 + UpstreamCredentialProvider, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 + to_server_spec, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( # noqa: PLC0415 + OAuthToken, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.presented_token_store import ( # noqa: PLC0415 + PresentedOAuthTokenStore, + ) + + forwarded_authorization = ( + effective_oauth2_headers.get("Authorization") + if effective_oauth2_headers + else None + ) + is_interactive_authz_code = ( + server_model.auth_type == MCPAuth.oauth2 + and forwarded_authorization is not None + and to_server_spec(server_model) is not None + ) + preview_cred_provider = ( + UpstreamCredentialProvider( + oauth_token_store=PresentedOAuthTokenStore( + OAuthToken( + access_token=forwarded_authorization[7:] + if forwarded_authorization[:7].lower() == "bearer " + else forwarded_authorization + ) + ) + ) + if is_interactive_authz_code + else None + ) + merged_headers = merge_mcp_headers( - extra_headers=effective_oauth2_headers, + extra_headers=( + None if preview_cred_provider else effective_oauth2_headers + ), static_headers=request.static_headers, ) @@ -1044,6 +1090,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, extra_headers=merged_headers, stdio_env=stdio_env, + cred_provider=preview_cred_provider, ) return await operation(client) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index d7975303802..55a0641c887 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -283,6 +283,7 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _should_strip_caller_authorization, + _without_authorization, global_mcp_server_manager, ) from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( @@ -1360,115 +1361,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], @@ -1529,6 +1436,16 @@ if MCP_AVAILABLE: else: # Copy to avoid mutating the original dict (important for parallel fetching) extra_headers = oauth2_headers.copy() if oauth2_headers else None + # 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 _call_regular_mcp_tool. + if extra_headers and _should_strip_caller_authorization( + mcp_server=server, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + ): + extra_headers = _without_authorization(extra_headers) if server.extra_headers and raw_headers: if extra_headers is None: @@ -1778,8 +1695,17 @@ if MCP_AVAILABLE: # Prefer server-stored per-user OAuth when configured, so a stale # Authorization header from the MCP client cannot override Redis/DB # (same issue as call_tool in mcp_server_manager: VS Code caches tokens). + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 + to_server_spec, + ) + + # A server migrated to the v2 resolver gets its token from the resolver at connect + # time; building it here would double-resolve and be shadowed by the v2 graft. The + # preemptive 401 already challenged a missing token, so one exists for the connect. + migrated_to_v2 = to_server_spec(server) is not None if ( - server.auth_type == MCPAuth.oauth2 + not migrated_to_v2 + and server.auth_type == MCPAuth.oauth2 and getattr(server, "needs_user_oauth_token", False) and user_api_key_auth is not None ): @@ -1792,7 +1718,11 @@ if MCP_AVAILABLE: extra_headers = db_headers # If still no OAuth2 token, fall back to pre-fetched creds (non-stale-client path) - elif extra_headers is None and server.auth_type == MCPAuth.oauth2: + elif ( + not migrated_to_v2 + and extra_headers is None + and server.auth_type == MCPAuth.oauth2 + ): extra_headers = await _get_user_oauth_extra_headers_from_db( server, user_api_key_auth, @@ -3501,12 +3431,6 @@ if MCP_AVAILABLE: # If no stored token exists, fail fast with 401 so clients can # kick off PKCE/interactive OAuth flow immediately. if server.needs_user_oauth_token: - stored_oauth_headers = await _get_user_oauth_extra_headers_from_db( - server=server, - user_api_key_auth=user_api_key_auth, - ) - if stored_oauth_headers: - continue if getattr(server, "delegate_auth_to_upstream", False) is True: # Delegate-auth servers run upstream PKCE: challenge with # the proxied resource_metadata (RFC 9728), not the @@ -3521,6 +3445,12 @@ if MCP_AVAILABLE: detail="Unauthorized", headers={"www-authenticate": www_authenticate}, ) + # The v2 resolver owns the existence check, so every authorization_code + # resolution (egress and this discovery challenge) runs through it. + if await global_mcp_server_manager.has_user_oauth_token( + server, user_api_key_auth + ): + continue request = StarletteRequest(scope) base_url = get_request_base_url(request) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 594e9dcc969..383dc255607 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -7,17 +7,20 @@ maps each CredError onto its HTTP status. These pin the parity-critical mapping import base64 from types import SimpleNamespace +from unittest.mock import patch import pytest from fastapi import HTTPException 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.types import ( ApiKeyConfig, + AuthorizationCodeConfig, CredError, NoneConfig, SharedKey, @@ -71,12 +74,27 @@ def test_basic_scheme_base64_encodes_the_token(): assert spec.config.key_source.value.get_secret_value() == expected +@pytest.mark.parametrize( + "oauth2_flow", + [None, "authorization_code"], +) +def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow): + # oauth2 without client_credentials is the per-user authorization_code mode. + spec = to_server_spec(_server(auth_type=MCPAuth.oauth2, oauth2_flow=oauth2_flow)) + assert spec is not None and isinstance(spec.config, AuthorizationCodeConfig) + + @pytest.mark.parametrize( "server", [ _server(auth_type=MCPAuth.api_key), # no token configured _server(auth_type=MCPAuth.bearer_token), # no token configured - _server(auth_type=MCPAuth.oauth2), + _server( + auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials" + ), # M2M -> v1 + _server( + auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True + ), # delegated upstream OAuth -> v1 _server(auth_type=MCPAuth.oauth2_token_exchange), _server(auth_type=MCPAuth.aws_sigv4), _server( @@ -162,3 +180,43 @@ def test_raise_public_plain_unauthorized_has_no_challenge(): assert exc.status_code == 401 assert exc.detail == "unauthorized: nope" assert exc.headers is None + + +_ROOT_PATH = "litellm.proxy.utils.get_server_root_path" + + +def test_raise_user_oauth_challenge_points_at_per_server_prm(): + with patch(_ROOT_PATH, return_value="/"), pytest.raises(HTTPException) as exc_info: + raise_user_oauth_challenge(_server(alias="my-srv")) + exc = exc_info.value + assert exc.status_code == 401 + assert ( + exc.headers["WWW-Authenticate"] + == 'Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/my-srv"' + ) + + +def test_raise_user_oauth_challenge_includes_server_root_path(): + with ( + patch(_ROOT_PATH, return_value="/api/v1"), + pytest.raises(HTTPException) as exc_info, + ): + raise_user_oauth_challenge(_server(alias="my-srv")) + assert ( + exc_info.value.headers["WWW-Authenticate"] + == 'Bearer resource_metadata="/.well-known/oauth-protected-resource/api/v1/mcp/my-srv"' + ) + + +@pytest.mark.parametrize( + "kwargs, expected_name", + [ + ({"alias": "a", "server_name": "sn"}, "a"), # alias wins + ({"server_name": "sn"}, "sn"), # then server_name + ({}, "n"), # then the name field (server_id is the last fallback) + ], +) +def test_raise_user_oauth_challenge_name_fallback(kwargs, expected_name): + with patch(_ROOT_PATH, return_value="/"), pytest.raises(HTTPException) as exc_info: + raise_user_oauth_challenge(_server(**kwargs)) + assert f'/mcp/{expected_name}"' in exc_info.value.headers["WWW-Authenticate"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py new file mode 100644 index 00000000000..91dd1aa5cc6 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py @@ -0,0 +1,184 @@ +"""Tests for the authorization_code refresher: the refresh_token grant, parsing, and persist.""" + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import ( + AuthorizationCodeRefresher, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, +) + + +class _Server: + def __init__( + self, + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="sec", + ): + self.token_url = token_url + self.client_id = client_id + self.client_secret = client_secret + + +def _lookup(server): + return lambda server_id: server + + +def _endpoint(body, sink=None): + async def post(url, form): + if sink is not None: + sink.append((url, form)) + return body + + return post + + +def _recording_persist(sink): + async def persist( + user_id, server_id, access_token, refresh_token, expires_in, scopes + ): + sink.append( + (user_id, server_id, access_token, refresh_token, expires_in, scopes) + ) + + return persist + + +def _refresher( + server=None, body=None, *, post_sink=None, persist_sink=None, clock=lambda: 1000.0 +): + return AuthorizationCodeRefresher( + _lookup(server if server is not None else _Server()), + _endpoint(body, post_sink), + _recording_persist(persist_sink if persist_sink is not None else []), + clock=clock, + ) + + +@pytest.mark.asyncio +async def test_refreshes_persists_and_returns_typed_token(): + persisted = [] + posted = [] + refresher = _refresher( + body={ + "access_token": "new-at", + "expires_in": 3600, + "refresh_token": "new-rt", + "scope": "a b", + }, + post_sink=posted, + persist_sink=persisted, + ) + token = await refresher.refresh( + "alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt") + ) + + assert token is not None + assert token.access_token == "new-at" + assert token.refresh_token == "new-rt" + assert token.expires_at == 1000.0 + 3600 # clock + expires_in -> epoch + # the rotated triple is persisted for (user, server) with parsed scopes + assert persisted == [("alice", "srv", "new-at", "new-rt", 3600, ("a", "b"))] + # the grant carried the refresh_token + client credentials + url, form = posted[0] + assert url == "https://idp.example.com/token" + assert form == { + "grant_type": "refresh_token", + "refresh_token": "old-rt", + "client_id": "cid", + "client_secret": "sec", + } + + +@pytest.mark.asyncio +async def test_no_refresh_token_is_not_refreshable(): + posted = [] + refresher = _refresher(body={"access_token": "x"}, post_sink=posted) + assert ( + await refresher.refresh("alice", "srv", OAuthToken(access_token="old")) is None + ) + assert posted == [] # never hit the IdP + + +@pytest.mark.asyncio +async def test_unknown_server_or_no_token_url_yields_none(): + assert ( + await _refresher(server=None).refresh( + "a", "s", OAuthToken("old", refresh_token="rt") + ) + is None + ) + no_url = _Server(token_url=None) + assert ( + await _refresher(server=no_url).refresh( + "a", "s", OAuthToken("old", refresh_token="rt") + ) + is None + ) + + +@pytest.mark.asyncio +async def test_grant_failure_does_not_persist(): + persisted = [] + refresher = _refresher( + body=None, persist_sink=persisted + ) # token_endpoint signals failure + assert ( + await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None + ) + assert persisted == [] + + +@pytest.mark.asyncio +async def test_response_without_access_token_does_not_persist(): + persisted = [] + refresher = _refresher(body={"expires_in": 60}, persist_sink=persisted) + assert ( + await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None + ) + assert persisted == [] + + +@pytest.mark.asyncio +async def test_unrotated_refresh_token_is_carried_forward(): + persisted = [] + refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted) + token = await refresher.refresh( + "a", "s", OAuthToken("old", refresh_token="keep-rt") + ) + assert token is not None + assert ( + token.refresh_token == "keep-rt" + ) # response omitted refresh_token -> reuse the old one + assert token.expires_at is None # no expires_in -> no known expiry + assert persisted[0][3] == "keep-rt" + + +@pytest.mark.asyncio +async def test_unrecorded_scope_is_carried_forward(): + persisted = [] + refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted) + token = await refresher.refresh( + "a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")) + ) + assert token is not None + # response omitted "scope" -> the user's recorded grant is preserved, not dropped + assert token.scopes == ("read", "write") + assert persisted[0][5] == ("read", "write") + + +@pytest.mark.asyncio +async def test_returned_scope_overrides_prior_when_present(): + persisted = [] + refresher = _refresher( + body={"access_token": "new-at", "scope": "read"}, + persist_sink=persisted, + ) + token = await refresher.refresh( + "a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write")) + ) + assert token is not None + assert token.scopes == ("read",) # a present scope replaces the prior grant + assert persisted[0][5] == ("read",) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_presented_token_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_presented_token_store.py new file mode 100644 index 00000000000..8b715fa717b --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_presented_token_store.py @@ -0,0 +1,19 @@ +"""Tests for the create/test-preview presented OAuth token store.""" + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.presented_token_store import ( + PresentedOAuthTokenStore, +) + + +@pytest.mark.asyncio +async def test_serves_the_presented_token_regardless_of_key(): + token = OAuthToken(access_token="at", scopes=("read",)) + store = PresentedOAuthTokenStore(token) + # one-shot store backs a single preview call, so the lookup key is irrelevant + assert await store.fetch("alice", "srv-1") is token + assert await store.fetch("", "") is token diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 75be6dfc157..73e9a52b937 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -1,9 +1,9 @@ """Tests for the resolver dispatch: live arms produce auth, stubbed arms fail closed. -`none` and `api_key` (shared-key source) are implemented; every other arm, plus the `api_key` -BYOK source, returns a typed `not_implemented` error until its mode lands. Parametrizing the -stubs over one config each also guards reachability: a dropped `case` would hit `assert_never` -and raise instead of returning the stub. +`none`, `api_key` (shared-key source), and `authorization_code` are implemented; every other arm, +plus the `api_key` BYOK source, returns a typed `not_implemented` error until its mode lands. +Parametrizing the stubs over one config each also guards reachability: a dropped `case` would hit +`assert_never` and raise instead of returning the stub. """ import httpx @@ -28,6 +28,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( TokenExchangeConfig, UpstreamCredentialProvider, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, + TokenStoreUnavailable, +) _SUBJECT = Subject(tenant_id="", subject_id="") @@ -84,12 +88,114 @@ async def test_api_key_shared_honors_authorization_scheme(): assert _emitted(result.ok)["Authorization"] == "Bearer tok" +class _FakeTokenStore: + """An OAuthTokenStore returning a canned per-user token (None == not authorized).""" + + def __init__(self, by_user: dict) -> None: + self._by_user = by_user + + async def fetch(self, user_id: str, server_id: str): + return self._by_user.get((user_id, server_id)) + + +@pytest.mark.asyncio +async def test_authorization_code_emits_bearer_for_a_stored_token(): + store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="at-alice")}) + result = await UpstreamCredentialProvider( + oauth_token_store=store + ).resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) + ) + assert isinstance(result, Ok) + assert _emitted(result.ok)["Authorization"] == "Bearer at-alice" + + +@pytest.mark.asyncio +async def test_authorization_code_without_token_is_semantically_unauthorized(): + result = await UpstreamCredentialProvider( + oauth_token_store=_FakeTokenStore({}) + ).resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) + ) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + # Semantic only: the per-server challenge is built at the edge, not in the request-free arm. + assert "Authorization required" in result.error.unauthorized.detail + assert result.error.unauthorized.www_authenticate is None + assert result.error.unauthorized.body is None + + +@pytest.mark.asyncio +async def test_authorization_code_store_unavailable_is_unauthorized(): + class _Unavailable: + async def fetch(self, user_id: str, server_id: str): + raise TokenStoreUnavailable("down") + + result = await UpstreamCredentialProvider( + oauth_token_store=_Unavailable() + ).resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) + ) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + + +@pytest.mark.asyncio +async def test_authorization_code_with_no_store_wired_is_unauthorized(): + result = await UpstreamCredentialProvider().resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) + ) + assert isinstance(result, Error) + assert result.error.tag == "unauthorized" + + +@pytest.mark.asyncio +async def test_authorization_code_isolates_by_subject(): + store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="at-alice")}) + provider = UpstreamCredentialProvider(oauth_token_store=store) + alice = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig()) + ) + bob = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="bob"), _spec(AuthorizationCodeConfig()) + ) + assert ( + isinstance(alice, Ok) + and _emitted(alice.ok)["Authorization"] == "Bearer at-alice" + ) + assert isinstance(bob, Error) and bob.error.tag == "unauthorized" + + +@pytest.mark.asyncio +async def test_has_user_token_reflects_the_stored_token(): + present = UpstreamCredentialProvider( + oauth_token_store=_FakeTokenStore( + {("alice", "s"): OAuthToken(access_token="at")} + ) + ) + absent = UpstreamCredentialProvider(oauth_token_store=_FakeTokenStore({})) + spec = _spec(AuthorizationCodeConfig()) + subject = Subject(tenant_id="", subject_id="alice") + assert await present.has_user_token(subject, spec) is True + assert await absent.has_user_token(subject, spec) is False + + +@pytest.mark.asyncio +async def test_has_user_token_false_for_a_non_per_user_mode(): + # A none-mode server has no per-user token to check. + provider = UpstreamCredentialProvider() + spec = _spec(NoneConfig()) + assert ( + await provider.has_user_token(Subject(tenant_id="", subject_id="a"), spec) + is False + ) + + _STUBBED = [ ("api_key_byok", ApiKeyConfig(key_source=Byok())), ("passthrough", PassthroughConfig()), ("client_credentials", ClientCredentialsConfig()), ("token_exchange", TokenExchangeConfig()), - ("authorization_code", AuthorizationCodeConfig()), ("aws_sigv4", AwsSigV4Config(region="us-east-1")), ] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v2_token_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v2_token_store.py new file mode 100644 index 00000000000..97d14b65f5a --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_v2_token_store.py @@ -0,0 +1,73 @@ +"""Tests for the v2-native per-user token read store: validate the decoded blob into a typed token.""" + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import ( + V2PerUserTokenStore, +) + + +def _reader(payload): + async def read(user_id: str, server_id: str): + return payload + + return read + + +@pytest.mark.asyncio +async def test_builds_typed_token_with_epoch_expiry(): + store = V2PerUserTokenStore( + _reader( + { + "access_token": "at", + "expires_at": "2099-01-01T00:00:00+00:00", + "refresh_token": "rt", + } + ) + ) + token = await store.fetch("alice", "srv") + assert token is not None + assert token.access_token == "at" + assert token.refresh_token == "rt" + # ISO string is converted to epoch seconds for RefreshingTokenStore to compare against. + assert token.expires_at == 4070908800.0 + + +@pytest.mark.asyncio +async def test_optional_fields_absent_yield_none(): + store = V2PerUserTokenStore(_reader({"access_token": "at"})) + token = await store.fetch("alice", "srv") + assert token is not None + assert token.expires_at is None and token.refresh_token is None + + +@pytest.mark.asyncio +async def test_unparseable_expiry_is_dropped_not_raised(): + store = V2PerUserTokenStore( + _reader({"access_token": "at", "expires_at": "not-a-date"}) + ) + token = await store.fetch("alice", "srv") + assert token is not None and token.expires_at is None + + +@pytest.mark.asyncio +async def test_missing_access_token_is_not_authorized(): + store = V2PerUserTokenStore(_reader({"expires_at": "2099-01-01T00:00:00+00:00"})) + assert await store.fetch("alice", "srv") is None + + +@pytest.mark.asyncio +async def test_no_credential_is_not_authorized(): + assert await V2PerUserTokenStore(_reader(None)).fetch("alice", "srv") is None + + +@pytest.mark.asyncio +async def test_empty_user_short_circuits_without_reading(): + calls = [] + + async def read(user_id: str, server_id: str): + calls.append((user_id, server_id)) + return {"access_token": "at"} + + assert await V2PerUserTokenStore(read).fetch("", "srv") is None + assert calls == [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index ca326d197ae..34e932b6ae7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -284,8 +284,11 @@ def test_prepare_mcp_server_headers_oauth2_m2m_omits_litellm_caller_authorizatio assert extra_headers is None -def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers(): - """Interactive OAuth still forwards the user's OAuth token in extra_headers.""" +def test_prepare_mcp_server_headers_oauth2_interactive_drops_caller_authorization(): + """A v2-migrated interactive OAuth (authorization_code) server must NOT forward the + caller's Authorization: the resolver injects the stored per-user token, so a + caller-supplied bearer must not override another user's stored credential. Non-auth + headers are still carried; only the credential is dropped.""" try: from litellm.proxy._experimental.mcp_server.server import ( _prepare_mcp_server_headers, @@ -293,7 +296,7 @@ def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers(): except ImportError: pytest.skip("MCP server not available") - user_oauth = {"Authorization": "Bearer upstream-user-token"} + caller_oauth = {"Authorization": "Bearer caller-supplied-token"} server = MCPServer( server_id="3lo-server", @@ -307,12 +310,13 @@ def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers(): server=server, mcp_server_auth_headers=None, mcp_auth_header=None, - oauth2_headers=user_oauth, + oauth2_headers=caller_oauth, raw_headers=None, ) assert server_auth_header is None - assert extra_headers == user_oauth + # Caller's Authorization is dropped (only key present) -> extra_headers is None. + assert extra_headers is None def test_prepare_mcp_server_headers_m2m_skips_authorization_from_raw_extra_headers(): @@ -2813,8 +2817,10 @@ async def test_mcp_routing_with_conflicting_alias_and_group_name(): @pytest.mark.asyncio @pytest.mark.no_parallel -async def test_oauth2_headers_passed_to_mcp_client(): - """Test that OAuth2 headers are properly passed through to the MCP client for OAuth2 servers like github_mcp""" +async def test_oauth2_caller_headers_not_forwarded_for_migrated_server(): + """A v2-migrated authorization_code server (like github_mcp) must NOT forward the + caller's oauth2 Authorization to the MCP client — the resolver injects the stored + per-user token, so a caller-supplied bearer cannot override another user's credential.""" try: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -2928,20 +2934,13 @@ async def test_oauth2_headers_passed_to_mcp_client(): assert captured_client_args["server"].server_id == oauth2_server.server_id assert captured_client_args["server"].auth_type == MCPAuth.oauth2 - # Most importantly: verify that OAuth2 headers were passed as extra_headers - assert ( - captured_client_args["extra_headers"] is not None - ), "Expected extra_headers to be passed for OAuth2 server" - assert ( - captured_client_args["extra_headers"] == oauth2_headers - ), f"Expected OAuth2 headers to be passed as extra_headers, got {captured_client_args['extra_headers']}" - - # Verify the Authorization header specifically - assert "Authorization" in captured_client_args["extra_headers"] - assert ( - captured_client_args["extra_headers"]["Authorization"] - == "Bearer github_oauth_token_12345" - ) + # Security: a v2-migrated authorization_code server must NOT forward the caller's + # oauth2 Authorization upstream. The v2 resolver injects the stored per-user token, + # so a caller-supplied bearer cannot override another user's stored credential. + extra_headers = captured_client_args["extra_headers"] + assert extra_headers is None or "Authorization" not in { + k.lower() for k in extra_headers + }, f"Caller Authorization must not be forwarded, got {extra_headers}" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 35a67391315..e6c2b57ee79 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -31,6 +31,8 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _deserialize_json_dict, _deserialize_json_list, _normalize_mcp_server_cost_info, + _should_strip_caller_authorization, + _without_authorization, ) from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -137,6 +139,43 @@ class TestMCPServerManager: assert client.stdio_config["env"]["NODE_ENV"] == "test" assert client.stdio_config["env"]["NPM_CONFIG_CACHE"] == MCP_NPM_CACHE_DIR + @pytest.mark.asyncio + async def test_caller_auth_header_cannot_bypass_v2_for_authorization_code(self): + """A caller-supplied per-request override must not substitute the stored authorization_code + token: _create_mcp_client keeps the v2 spec and resolves through the injected provider + rather than deferring to the v1 caller-override path.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Ok, + ) + from litellm.types.mcp import MCPAuth + + calls = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + calls.append((subject.subject_id, server.server_id)) + return Ok(StaticHeaderAuth("stored-token")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = MCPServer( + server_id="authz-srv", + name="authz", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2, # oauth2 + no client creds + not delegate -> authorization_code + ) + + client = await manager._create_mcp_client( + server, mcp_auth_header="Bearer caller-supplied-token" + ) + + # the v2 resolver ran (the caller override did NOT defer to v1); the stored token wins + assert calls == [("", "authz-srv")] + assert client is not None + async def test_create_mcp_client_stdio_injects_npm_config_cache(self): """Test that _create_mcp_client injects NPM_CONFIG_CACHE when not already set, and preserves user-provided NPM_CONFIG_CACHE when present.""" @@ -600,6 +639,84 @@ class TestMCPServerManager: assert captured_extra_headers == {"x-request-id": "req-123"} + @pytest.mark.asyncio + async def test_call_regular_mcp_tool_v2_authz_code_drops_caller_authorization( + self, + ): + """A v2-migrated per-user OAuth (authorization_code) server must NOT seed a + caller-forwarded Authorization into extra_headers — the resolver injects the + stored per-user token, and apply-if-absent would otherwise let the caller's + header override another user's stored credential (matches v1's overwrite).""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + # oauth2, not M2M, not delegate => to_server_spec maps it to AuthorizationCodeConfig + server = MCPServer( + server_id="server-authz-code-call", + name="authz-code-server", + url="https://example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + # Migrated authorization_code => the centralized strip decision says drop the + # caller's Authorization (the v2 resolver injects the stored token). + assert ( + _should_strip_caller_authorization( + mcp_server=server, raw_headers=None, user_api_key_auth=None + ) + is True + ) + + mock_client = AsyncMock() + mock_client.call_tool = AsyncMock( + return_value=CallToolResult(content=[], isError=False) + ) + captured_extra_headers = "unset" + + async def capture_create_mcp_client( + server, + mcp_auth_header, + extra_headers, + stdio_env, + subject_token=None, + **kwargs, + ): # pragma: no cover - helper + nonlocal captured_extra_headers + captured_extra_headers = extra_headers + return mock_client + + manager._create_mcp_client = AsyncMock(side_effect=capture_create_mcp_client) + + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="tool", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer caller-supplied-token"}, + raw_headers={"authorization": "Bearer caller-supplied-token"}, + proxy_logging_obj=None, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), + ) + + # The caller's Authorization must not reach extra_headers; the v2 resolver is + # the sole Authorization source for this server. + assert captured_extra_headers != "unset" + if captured_extra_headers: + assert "authorization" not in {k.lower() for k in captured_extra_headers} + + def test_without_authorization_drops_only_the_credential(self): + # None / empty -> None + assert _without_authorization(None) is None + assert _without_authorization({}) is None + # Only Authorization present -> nothing left -> None (case-insensitive) + assert _without_authorization({"authorization": "Bearer x"}) is None + # Authorization dropped, other headers kept + assert _without_authorization( + {"Authorization": "Bearer x", "X-Trace-Id": "t"} + ) == {"X-Trace-Id": "t"} + @pytest.mark.asyncio async def test_call_regular_mcp_tool_passthrough_forwards_authorization_with_admission_header( self, @@ -2221,7 +2338,7 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_resolve_oauth2_headers_looks_up_stored_token(self): - """Falls back to stored per-user OAuth headers when no token is supplied.""" + """Falls back to stored per-user OAuth headers for a non-migrated (delegate) oauth2 server.""" from litellm.proxy._types import UserAPIKeyAuth manager = MCPServerManager() @@ -2230,6 +2347,7 @@ class TestMCPServerManager: name="oauth-srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, # stays on v1, so v1 still builds the header ) user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") stored = {"Authorization": "Bearer stored-user-token"} @@ -2245,9 +2363,36 @@ class TestMCPServerManager: assert result == stored mock_lookup.assert_awaited_once() + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_steps_aside_for_migrated_server(self): + """A migrated authorization_code server is owned by the v2 resolver, so v1 must not also + build the token into extra_headers (which the v2 graft would defer to and shadow). + """ + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, # per-user, not delegate -> migrated to v2 + ) + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(return_value={"Authorization": "Bearer should-not-be-used"}), + ) as mock_lookup: + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + + assert result is None # stepped aside; the v2 resolver handles the token + mock_lookup.assert_not_awaited() + @pytest.mark.asyncio async def test_resolve_oauth2_headers_swallows_lookup_exception(self): - """Returns supplied headers (None) when the stored-token lookup raises.""" + """Returns supplied headers (None) when the v1 stored-token lookup raises (delegate path).""" from litellm.proxy._types import UserAPIKeyAuth manager = MCPServerManager() @@ -2256,6 +2401,7 @@ class TestMCPServerManager: name="oauth-srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, # non-migrated, so it reaches the v1 lookup ) user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") @@ -2268,6 +2414,51 @@ class TestMCPServerManager: ) assert result is None + @pytest.mark.asyncio + async def test_has_user_oauth_token_delegates_to_provider(self): + """has_user_oauth_token maps the server and delegates the verdict to the v2 resolver.""" + from litellm.proxy._types import UserAPIKeyAuth + + for verdict in (True, False): + + class _Provider: + async def has_user_token(self, subject, spec): + return verdict + + manager = MCPServerManager(cred_provider=_Provider()) + server = MCPServer( + server_id="s", + name="n", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + user_auth = UserAPIKeyAuth(api_key="sk", user_id="alice") + assert await manager.has_user_oauth_token(server, user_auth) is verdict + + @pytest.mark.asyncio + async def test_has_user_oauth_token_short_circuits_for_unmigrated_server(self): + """A server the resolver does not own (None spec, e.g. delegate) is False without a call.""" + from litellm.proxy._types import UserAPIKeyAuth + + calls: list = [] + + class _Provider: + async def has_user_token(self, subject, spec): + calls.append(spec) + return True + + manager = MCPServerManager(cred_provider=_Provider()) + server = MCPServer( + server_id="s", + name="n", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + ) + user_auth = UserAPIKeyAuth(api_key="sk", user_id="alice") + assert await manager.has_user_oauth_token(server, user_auth) is False + assert calls == [] # short-circuited on the None spec, never hit the resolver + @pytest.mark.asyncio async def test_resolve_oauth2_headers_no_user_id(self): """Skip lookup entirely when user_api_key_auth has no user_id.""" @@ -2356,7 +2547,6 @@ class TestMCPServerManager: # Unprefixed resolution resolved_server_unpref = manager._get_mcp_server_from_tool_name("create_zap") - print(resolved_server_unpref) assert resolved_server_unpref is not None assert resolved_server_unpref.server_id == server.server_id @@ -2970,14 +3160,17 @@ class TestMCPServerManager: object_permission_id="perm_no_mcp", ) - with patch.object( - manager, "get_allow_all_keys_server_ids", return_value=["global-server"] - ), patch.object( - MCPRequestHandler, - "get_allowed_mcp_servers", - new_callable=AsyncMock, - return_value=["leaked-server"], - ) as mock_inner: + with ( + patch.object( + manager, "get_allow_all_keys_server_ids", return_value=["global-server"] + ), + patch.object( + MCPRequestHandler, + "get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=["leaked-server"], + ) as mock_inner, + ): result = await manager.get_allowed_mcp_servers(user_api_key_auth) assert result == [] @@ -4641,15 +4834,23 @@ class TestCreateMcpClientV2Graft: client._resolved_auth._header_value.get_secret_value() == f"Basic {encoded}" ) - async def test_deferred_mode_uses_v1_auth_value(self): + async def test_m2m_client_credentials_defers_to_v1(self): + # M2M (oauth2 client_credentials) is not migrated: to_server_spec returns + # None, so the graft sets no resolved auth and leaves v1 in charge (v1 + # performs the client_credentials grant itself - the static + # authentication_token is never consumed for oauth2, so it does not flow + # to _mcp_auth_value). Per-user oauth2 (authorization_code) is migrated to + # v2 and is exercised separately. client = await MCPServerManager()._create_mcp_client( self._http_server( - auth_type=MCPAuth.oauth2, authentication_token="legacy-token" + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + authentication_token="legacy-token", ) ) assert client._resolved_auth is None - assert client._mcp_auth_value == "legacy-token" + assert client._mcp_auth_value is None async def test_static_token_missing_defers_to_v1(self): client = await MCPServerManager()._create_mcp_client( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 67b72dbb695..e4e1890a45a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -628,6 +628,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): oauth_server = MagicMock() oauth_server.auth_type = MCPAuth.oauth2 oauth_server.needs_user_oauth_token = True + oauth_server.delegate_auth_to_upstream = False with ( patch( @@ -648,10 +649,10 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token", new_callable=AsyncMock, - return_value=None, - ) as mock_get_stored_token, + return_value=False, + ) as mock_has_token, patch( "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", return_value=oauth_server, @@ -666,7 +667,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): await handle_streamable_http_mcp(scope, receive, send) # Verify a 401 was raised - assert mock_get_stored_token.await_count == 1 + assert mock_has_token.await_count == 1 assert mock_handle_request.await_count == 0 assert exc_info.value.status_code == 401 assert "www-authenticate" in exc_info.value.headers @@ -817,6 +818,7 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): oauth_server = MagicMock() oauth_server.auth_type = MCPAuth.oauth2 oauth_server.needs_user_oauth_token = True + oauth_server.delegate_auth_to_upstream = False with ( patch( @@ -837,10 +839,10 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token", new_callable=AsyncMock, - return_value={"Authorization": "Bearer cached-token"}, - ) as mock_get_stored_token, + return_value=True, + ) as mock_has_token, patch( "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", return_value=oauth_server, @@ -858,7 +860,7 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): ): await handle_streamable_http_mcp(scope, receive, send) - assert mock_get_stored_token.await_count == 1 + assert mock_has_token.await_count == 1 assert mock_handle_request.await_count == 1 @@ -941,10 +943,9 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token", new_callable=AsyncMock, - return_value=None, - ) as mock_get_stored_token, + ) as mock_has_token, patch( "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", return_value=delegated_server, @@ -958,7 +959,9 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns with pytest.raises(HTTPException) as exc_info: await handle_streamable_http_mcp(scope, receive, send) - assert mock_get_stored_token.await_count == 1 + # Delegate-auth servers raise the resource_metadata challenge before any + # per-user existence check, so the v2 token store is never consulted. + assert mock_has_token.await_count == 0 assert mock_handle_request.await_count == 0 assert exc_info.value.status_code == 401 challenge = exc_info.value.headers["www-authenticate"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 27f9a311250..9e3862b43eb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -270,6 +270,155 @@ class TestExecuteWithMcpClient: or "Authorization" not in captured["extra_headers"] ) + @pytest.mark.asyncio + async def test_interactive_oauth_resolves_forwarded_token_via_presented_store( + self, monkeypatch + ): + """Interactive authorization_code preview (oauth2, no client credentials): the forwarded + just-authorized token is resolved THROUGH the v2 resolver via a one-shot presented store + (cred_provider), not the caller-override path. The bare token (Bearer stripped) is the + upstream credential and is not also forwarded in extra_headers.""" + captured: dict = {} + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + captured.update(kwargs) + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="linear", + url="https://mcp.linear.app/mcp", + auth_type=MCPAuth.oauth2, + authorization_url="https://mcp.linear.app/authorize", + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, + ok_operation, + oauth2_headers={"Authorization": "Bearer forwarded-user-token"}, + ) + + assert result["status"] == "ok" + # Resolved via the v2 resolver, never the caller-override header + assert captured["mcp_auth_header"] is None + provider = captured["cred_provider"] + assert provider is not None + token = await provider._oauth_token_store.fetch("u", "s") + assert token is not None and token.access_token == "forwarded-user-token" + # The resolver supplies the bearer, so it is not also forwarded as a caller header + extra_headers = captured.get("extra_headers") or {} + assert not any(k.lower() == "authorization" for k in extra_headers) + + @pytest.mark.asyncio + async def test_m2m_does_not_build_presented_store(self, monkeypatch): + """M2M (client_credentials): to_server_spec returns None, so no presented provider is built; + the auto-fetch path is unchanged (no cred_provider, the incoming header dropped as before).""" + captured: dict = {} + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + captured.update(kwargs) + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="m2m-server", + url="https://example.com", + auth_type=MCPAuth.oauth2, + token_url="https://auth.example.com/token", + credentials={"client_id": "my-id", "client_secret": "my-secret"}, + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, + ok_operation, + oauth2_headers={"Authorization": "Bearer sk-litellm-api-key"}, + ) + + assert result["status"] == "ok" + assert captured.get("cred_provider") is None + assert captured["mcp_auth_header"] is None + + @pytest.mark.asyncio + async def test_token_exchange_does_not_build_presented_store(self, monkeypatch): + """OBO / token-exchange (auth_type oauth2_token_exchange, not oauth2): excluded by the + auth_type == oauth2 guard, so no presented provider is built and the v1 exchange path runs.""" + captured: dict = {} + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + captured.update(kwargs) + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="obo-server", + url="https://example.com", + auth_type=MCPAuth.oauth2_token_exchange, + token_url="https://auth.example.com/token", + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, + ok_operation, + oauth2_headers={"Authorization": "Bearer subject-jwt"}, + ) + + assert result["status"] == "ok" + assert captured.get("cred_provider") is None + @pytest.mark.asyncio async def test_catches_exception_group(self, monkeypatch): """MCP SDK's anyio TaskGroup raises BaseExceptionGroup which does not @@ -1660,9 +1809,9 @@ class TestPreviewOpenAPITools: names = [t["name"] for t in result["tools"]] anthropic_re = re.compile(r"^[a-zA-Z0-9_-]{1,128}$") for name in names: - assert anthropic_re.match( - name - ), f"preview tool name {name!r} violates ^[a-zA-Z0-9_-]+$" + assert anthropic_re.match(name), ( + f"preview tool name {name!r} violates ^[a-zA-Z0-9_-]+$" + ) assert "actions_download-job-logs-for-workflow-run" in names assert "pulls_list-files" in names @@ -1719,9 +1868,7 @@ class TestPreviewOpenAPITools: registered_summary_to_name: dict = {} - def fake_create_tool_function( - path, method, operation, base_url - ): # noqa: ANN001 + def fake_create_tool_function(path, method, operation, base_url): # noqa: ANN001 def _f(): return None @@ -1734,9 +1881,7 @@ class TestPreviewOpenAPITools: ) class _StubRegistry: - def register_tool( - self, name, description, input_schema, handler - ): # noqa: ANN001 + def register_tool(self, name, description, input_schema, handler): # noqa: ANN001 registered_summary_to_name[description] = name monkeypatch.setattr(