ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(mcp): add OBO MCP Auth (#27421)

Browse Source 
* feat(mcp): add oauth2 token exchange auth

Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>

* fix(mcp): cache token exchange fallback

Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>

---------

Co-authored-by: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com>
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
ishaan-berri 2026-05-07 15:35:21 -07:00 committed by GitHub
parent 3b78a3a545
commit e4c14862fc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 868 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
litellm/constants.py
View File

@ -161,6 +161,11 @@ MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset(
| (set(_MCP_STDIO_EXTRA_COMMANDS.split(",")) - {""})
)
# MCP OAuth2 Token Exchange (OBO) Defaults
MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE = int(
os.getenv("MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE", "500")
)
LITELLM_UI_ALLOW_HEADERS = [
"x-litellm-semantic-filter",
"x-litellm-semantic-filter-tools",

2
litellm/experimental_mcp_client/client.py
View File

@ -366,6 +366,8 @@ class MCPClient:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.token:
headers["Authorization"] = f"token {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.oauth2_token_exchange:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request

196
litellm/proxy/_experimental/mcp_server/auth/token_exchange.py Normal file
View File

@ -0,0 +1,196 @@
"""
OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers.
Exchanges a user's incoming JWT (subject_token) for a scoped access token
at an IDP's token exchange endpoint. The exchanged token is then used to
authenticate requests to the upstream MCP server.
See: https://datatracker.ietf.org/doc/html/rfc8693
"""
import asyncio
import hashlib
import weakref
from typing import TYPE_CHECKING, Dict, Tuple
import httpx
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
# RFC 8693 grant type constant
TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
DEFAULT_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token"
class TokenExchangeHandler:
"""Handles OAuth 2.0 Token Exchange (RFC 8693) for MCP servers.
Caches exchanged tokens keyed by ``hash(subject_token + server_id)`` so
repeated calls with the same user token skip the IDP round-trip.
"""
def __init__(self) -> None:
self._cache = InMemoryCache(
max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
)
# WeakValueDictionary so locks are GC'd once no coroutine holds a reference,
# preventing unbounded growth with many rotating user tokens.
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = (
weakref.WeakValueDictionary()
)
def _get_lock(self, cache_key: str) -> asyncio.Lock:
lock = self._locks.get(cache_key)
if lock is None:
lock = asyncio.Lock()
self._locks[cache_key] = lock
return lock
@staticmethod
def _cache_key(subject_token: str, server_id: str) -> str:
raw = f"{subject_token}:{server_id}"
return hashlib.sha256(raw.encode()).hexdigest()
async def exchange_token(
self,
subject_token: str,
server: "MCPServer",
) -> str:
"""Exchange *subject_token* for a scoped access token.
Returns the exchanged ``access_token`` string (suitable for a
``Bearer`` header).
Raises ``ValueError`` on configuration or IDP errors.
"""
cache_key = self._cache_key(subject_token, server.server_id)
# Fast path
cached = self._cache.get_cache(cache_key)
if cached is not None:
return cached
# Slow path — one exchange at a time per (user, server) pair
async with self._get_lock(cache_key):
cached = self._cache.get_cache(cache_key)
if cached is not None:
return cached
token, ttl = await self._do_exchange(subject_token, server)
self._cache.set_cache(cache_key, token, ttl=ttl)
return token
async def _do_exchange(
self,
subject_token: str,
server: "MCPServer",
) -> Tuple[str, int]:
"""POST to the token exchange endpoint with RFC 8693 parameters.
Returns ``(access_token, ttl_seconds)``.
"""
endpoint = server.token_exchange_endpoint or server.token_url
if not endpoint:
raise ValueError(
f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange "
f"but no token_exchange_endpoint or token_url configured"
)
if not server.client_id or not server.client_secret:
raise ValueError(
f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange "
f"but missing client_id or client_secret"
)
data: Dict[str, str] = {
"grant_type": TOKEN_EXCHANGE_GRANT_TYPE,
"subject_token": subject_token,
"subject_token_type": server.subject_token_type
or DEFAULT_SUBJECT_TOKEN_TYPE,
"client_id": server.client_id,
"client_secret": server.client_secret,
}
if server.audience:
data["audience"] = server.audience
if server.scopes:
data["scope"] = " ".join(server.scopes)
verbose_logger.debug(
"Exchanging token for MCP server %s at %s (audience=%s)",
server.server_id,
endpoint,
server.audience,
)
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
try:
response = await client.post(endpoint, data=data)
response.raise_for_status()
except httpx.HTTPStatusError as exc:
verbose_logger.debug(
"Token exchange IDP error for MCP server %s (status %d)",
server.server_id,
exc.response.status_code,
)
raise ValueError(
f"Token exchange for MCP server '{server.server_id}' "
f"failed with status {exc.response.status_code}"
) from exc
body = response.json()
if not isinstance(body, dict):
raise ValueError(
f"Token exchange response for MCP server '{server.server_id}' "
f"returned non-object JSON (got {type(body).__name__})"
)
access_token = body.get("access_token")
if not access_token:
raise ValueError(
f"Token exchange response for MCP server '{server.server_id}' "
f"missing 'access_token'"
)
raw_expires_in = body.get("expires_in")
try:
expires_in = (
int(raw_expires_in)
if raw_expires_in is not None
else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
)
except (TypeError, ValueError):
expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
ttl = max(
expires_in - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
)
verbose_logger.info(
"Token exchange succeeded for MCP server %s (expires in %ds)",
server.server_id,
expires_in,
)
return access_token, ttl
def invalidate(self, subject_token: str, server_id: str) -> None:
"""Remove a cached exchanged token (e.g. after a 401)."""
cache_key = self._cache_key(subject_token, server_id)
self._cache.delete_cache(cache_key)
# Module-level singleton
mcp_token_exchange_handler = TokenExchangeHandler()

62
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View File

@ -411,6 +411,15 @@ class MCPServerManager:
aws_role_name=server_config.get("aws_role_name", None),
aws_session_name=server_config.get("aws_session_name", None),
instructions=server_config.get("instructions", None),
# Token Exchange (OBO) fields
token_exchange_endpoint=server_config.get(
"token_exchange_endpoint", None
),
audience=server_config.get("audience", None),
subject_token_type=server_config.get(
"subject_token_type",
"urn:ietf:params:oauth:token-type:access_token",
),
)
self._assign_unique_short_prefix(new_server)
self.config_mcp_servers[server_id] = new_server
@ -765,6 +774,17 @@ class MCPServerManager:
aws_role_name=aws_creds.get("aws_role_name"),
aws_session_name=aws_creds.get("aws_session_name"),
instructions=mcp_server.instructions,
# Token Exchange (OBO) fields — read from credentials JSON blob
token_exchange_endpoint=(
credentials_dict.get("token_exchange_endpoint")
if credentials_dict
else None
),
audience=(credentials_dict.get("audience") if credentials_dict else None),
subject_token_type=(
credentials_dict.get("subject_token_type") if credentials_dict else None
)
or "urn:ietf:params:oauth:token-type:access_token",
)
return new_server
@ -1139,6 +1159,29 @@ class MCPServerManager:
#########################################################
# Methods that call the upstream MCP servers
#########################################################
@staticmethod
def _extract_bearer_token(
oauth2_headers: Optional[Dict[str, str]],
raw_headers: Optional[Dict[str, str]],
) -> Optional[str]:
"""Extract the bare Bearer token from oauth2_headers or raw_headers.
Returns the token string without the ``Bearer `` prefix, or ``None``
if no Authorization header is found.
"""
auth_value: Optional[str] = None
if oauth2_headers and "Authorization" in oauth2_headers:
auth_value = oauth2_headers["Authorization"]
elif raw_headers:
# raw_headers may have lowercase keys depending on the ASGI server
normalized = {k.lower(): v for k, v in raw_headers.items()}
auth_value = normalized.get("authorization")
if auth_value:
if auth_value.startswith("Bearer "):
return auth_value[len("Bearer ") :]
return auth_value
return None
def _build_stdio_env(
self,
server: MCPServer,
@ -1172,25 +1215,30 @@ class MCPServerManager:
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
extra_headers: Optional[Dict[str, str]] = None,
stdio_env: Optional[Dict[str, str]] = None,
subject_token: Optional[str] = None,
) -> MCPClient:
"""
Create an MCPClient instance for the given server.
Auth resolution (single place for all auth logic):
1. ``mcp_auth_header`` — per-request/per-user override
2. OAuth2 client_credentials token — auto-fetched and cached
3. ``server.authentication_token`` — static token from config/DB
2. OAuth2 Token Exchange (OBO) — exchange user token for scoped token
3. OAuth2 client_credentials token — auto-fetched and cached
4. ``server.authentication_token`` — static token from config/DB
Args:
server: The server configuration.
mcp_auth_header: Optional per-request auth override.
extra_headers: Additional headers to forward.
stdio_env: Environment variables for stdio transport.
subject_token: Optional user JWT for token exchange (OBO) flow.
Returns:
Configured MCP client instance.
"""
auth_value = await resolve_mcp_auth(server, mcp_auth_header)
auth_value = await resolve_mcp_auth(
server, mcp_auth_header, subject_token=subject_token
)
transport = server.transport or MCPTransport.sse
@ -2542,9 +2590,12 @@ class MCPServerManager:
if server_auth_header is None:
server_auth_header = mcp_auth_header
# oauth2 headers
# Extract subject token for OAuth2 Token Exchange (OBO) flow
subject_token: Optional[str] = None
extra_headers: Optional[Dict[str, str]] = None
if mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.auth_type == MCPAuth.oauth2_token_exchange:
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
elif mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.has_client_credentials:
# For M2M OAuth servers, Authorization must come from token fetch.
extra_headers = None
@ -2612,6 +2663,7 @@ class MCPServerManager:
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
)
call_tool_params = MCPCallToolRequestParams(

42
litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py
View File

@ -26,6 +26,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
from litellm.proxy._experimental.mcp_server.auth import token_exchange
from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
@ -50,12 +51,23 @@ class MCPOAuth2TokenCache(InMemoryCache):
def _get_lock(self, server_id: str) -> asyncio.Lock:
return self._locks.setdefault(server_id, asyncio.Lock())
async def async_get_token(self, server: "MCPServer") -> Optional[str]:
@staticmethod
def _has_client_credentials_config(server: "MCPServer") -> bool:
return bool(server.client_id and server.client_secret and server.token_url)
async def async_get_token(
self,
server: "MCPServer",
*,
require_client_credentials_flow: bool = True,
) -> Optional[str]:
"""Return a valid access token, fetching or refreshing as needed.
Returns ``None`` when the server lacks client credentials config.
"""
if not server.has_client_credentials:
if require_client_credentials_flow and not server.has_client_credentials:
return None
if not self._has_client_credentials_config(server):
return None
server_id = server.server_id
@ -263,16 +275,38 @@ mcp_per_user_token_cache = MCPPerUserTokenCache()
async def resolve_mcp_auth(
server: "MCPServer",
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
subject_token: Optional[str] = None,
) -> Optional[Union[str, Dict[str, str]]]:
"""Resolve the auth value for an MCP server.
Priority:
1. ``mcp_auth_header`` — per-request/per-user override
2. OAuth2 client_credentials token — auto-fetched and cached
3. ``server.authentication_token`` — static token from config/DB
2. OAuth2 Token Exchange (OBO / RFC 8693) — exchange user token for scoped token
3. OAuth2 client_credentials token — auto-fetched and cached
4. ``server.authentication_token`` — static token from config/DB
"""
if mcp_auth_header:
return mcp_auth_header
if server.has_token_exchange_config:
if subject_token:
return await token_exchange.mcp_token_exchange_handler.exchange_token(
subject_token, server
)
# No subject_token — fall back to client_credentials using the same client
# credentials and token_url so M2M scenarios still work.
if server.client_id and server.client_secret and server.token_url:
return await mcp_oauth2_token_cache.async_get_token(
server,
require_client_credentials_flow=False,
)
# OBO configured but no subject_token and missing client credentials — warn
# rather than silently proceeding unauthenticated.
verbose_logger.warning(
"MCP server '%s' is configured for token exchange (OBO) but no subject_token "
"was provided and client credentials (client_id/client_secret/token_url) are "
"incomplete. The request will proceed without authentication.",
server.server_id,
)
if server.has_client_credentials:
return await mcp_oauth2_token_cache.async_get_token(server)
return server.authentication_token

18
litellm/types/mcp.py
View File

@ -37,6 +37,7 @@ class MCPAuth(str, enum.Enum):
oauth2 = "oauth2"
aws_sigv4 = "aws_sigv4"
token = "token"
oauth2_token_exchange = "oauth2_token_exchange"
# MCP Literals
@ -54,6 +55,7 @@ MCPAuthType = Optional[
MCPAuth.oauth2,
MCPAuth.aws_sigv4,
MCPAuth.token,
MCPAuth.oauth2_token_exchange,
]
]
@ -117,6 +119,22 @@ class MCPCredentials(TypedDict, total=False):
aws_session_name: Optional[str]
"""Session name for STS AssumeRole (used in CloudTrail). Not a secret — stored unencrypted."""
audience: Optional[str]
"""
Target audience for OAuth 2.0 Token Exchange (RFC 8693)
"""
token_exchange_endpoint: Optional[str]
"""
IDP token endpoint for OAuth 2.0 Token Exchange (RFC 8693)
"""
subject_token_type: Optional[str]
"""
Subject token type for OAuth 2.0 Token Exchange (RFC 8693).
Default: urn:ietf:params:oauth:token-type:access_token
"""
class MCPServerCostInfo(TypedDict, total=False):
default_cost_per_query: Optional[float]

13
litellm/types/mcp_server/mcp_server_manager.py
View File

@ -57,6 +57,10 @@ class MCPServer(BaseModel):
aws_service_name: Optional[str] = None # defaults to "bedrock-agentcore"
aws_role_name: Optional[str] = None # IAM role ARN for STS AssumeRole
aws_session_name: Optional[str] = None # session name for CloudTrail auditing
# Token Exchange (OBO) fields — RFC 8693
token_exchange_endpoint: Optional[str] = None
audience: Optional[str] = None
subject_token_type: str = "urn:ietf:params:oauth:token-type:access_token"
# Stdio-specific fields
command: Optional[str] = None
args: Optional[List[str]] = None
@ -127,3 +131,12 @@ class MCPServer(BaseModel):
return any(h.lower() in auth_header_names for h in self.extra_headers)
return False
@property
def has_token_exchange_config(self) -> bool:
"""True if this server is configured for OAuth2 token exchange (OBO / RFC 8693)."""
return (
self.auth_type == MCPAuth.oauth2_token_exchange
and bool(self.client_id and self.client_secret)
and bool(self.token_exchange_endpoint or self.token_url)
)

511
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_token_exchange.py Normal file
View File

@ -0,0 +1,511 @@
"""
Tests for OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers.
Covers: exchange flow, caching, error handling, resolve_mcp_auth integration,
bearer token extraction, and config loading.
"""
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from litellm.proxy._experimental.mcp_server.auth.token_exchange import (
TOKEN_EXCHANGE_GRANT_TYPE,
TokenExchangeHandler,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
resolve_mcp_auth,
)
from litellm.proxy._types import LiteLLM_MCPServerTable, MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
def _obo_server(**overrides) -> MCPServer:
defaults = dict(
server_id="srv-obo-1",
name="test-obo",
url="https://mcp.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
client_id="litellm-client-id",
client_secret="litellm-client-secret",
token_exchange_endpoint="https://idp.example.com/oauth2/token",
audience="api://mcp-server",
scopes=["mcp.tools.read", "mcp.tools.execute"],
)
defaults.update(overrides)
return MCPServer(**defaults)
def _exchange_response(token="exchanged-tok-abc", expires_in=3600):
resp = MagicMock()
resp.json.return_value = {
"access_token": token,
"token_type": "Bearer",
"expires_in": expires_in,
}
resp.raise_for_status = MagicMock()
resp.text = ""
return resp
# ── Exchange Flow ──
@pytest.mark.asyncio
async def test_exchange_token_success():
"""Token exchange sends correct RFC 8693 parameters and returns access_token."""
handler = TokenExchangeHandler()
server = _obo_server()
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response("scoped-token-1")
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
result = await handler.exchange_token("user-jwt-xyz", server)
assert result == "scoped-token-1"
mock_client.post.assert_called_once()
_, kwargs = mock_client.post.call_args
data = kwargs["data"]
assert data["grant_type"] == TOKEN_EXCHANGE_GRANT_TYPE
assert data["subject_token"] == "user-jwt-xyz"
assert data["subject_token_type"] == "urn:ietf:params:oauth:token-type:access_token"
assert data["audience"] == "api://mcp-server"
assert data["scope"] == "mcp.tools.read mcp.tools.execute"
assert data["client_id"] == "litellm-client-id"
assert data["client_secret"] == "litellm-client-secret"
@pytest.mark.asyncio
async def test_exchange_token_no_audience():
"""When audience is None, it is omitted from the request."""
handler = TokenExchangeHandler()
server = _obo_server(audience=None)
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response()
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
await handler.exchange_token("user-jwt", server)
_, kwargs = mock_client.post.call_args
assert "audience" not in kwargs["data"]
@pytest.mark.asyncio
async def test_exchange_token_no_scopes():
"""When scopes is None, scope param is omitted from the request."""
handler = TokenExchangeHandler()
server = _obo_server(scopes=None)
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response()
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
await handler.exchange_token("user-jwt", server)
_, kwargs = mock_client.post.call_args
assert "scope" not in kwargs["data"]
# ── Caching ──
@pytest.mark.asyncio
async def test_exchange_token_cached():
"""Second call with same user token uses cache — only 1 HTTP POST."""
handler = TokenExchangeHandler()
server = _obo_server()
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response("cached-exchange-tok")
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
t1 = await handler.exchange_token("same-jwt", server)
t2 = await handler.exchange_token("same-jwt", server)
assert t1 == t2 == "cached-exchange-tok"
assert mock_client.post.call_count == 1
@pytest.mark.asyncio
async def test_different_user_tokens_not_shared():
"""Different user JWTs get different exchanged tokens."""
handler = TokenExchangeHandler()
server = _obo_server()
call_count = 0
async def mock_post(url, data=None):
nonlocal call_count
call_count += 1
resp = MagicMock()
resp.json.return_value = {
"access_token": f"exchanged-{call_count}",
"expires_in": 3600,
}
resp.raise_for_status = MagicMock()
return resp
mock_client = AsyncMock()
mock_client.post = mock_post
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
):
t1 = await handler.exchange_token("user-a-jwt", server)
t2 = await handler.exchange_token("user-b-jwt", server)
assert t1 == "exchanged-1"
assert t2 == "exchanged-2"
assert call_count == 2
# ── Error Handling ──
@pytest.mark.asyncio
async def test_exchange_token_http_error():
"""HTTP errors from the IDP are wrapped in a ValueError."""
handler = TokenExchangeHandler()
server = _obo_server()
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.text = "invalid_grant"
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
"Bad Request",
request=MagicMock(),
response=mock_response,
)
mock_client = AsyncMock()
mock_client.post.return_value = mock_response
with (
patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
),
pytest.raises(ValueError, match="failed with status 400"),
):
await handler.exchange_token("bad-jwt", server)
@pytest.mark.asyncio
async def test_exchange_token_http_error_does_not_log_response_body():
"""Raw IDP error bodies are not logged because they can contain credentials."""
handler = TokenExchangeHandler()
server = _obo_server()
raw_response_body = "client_secret=do-not-log"
mock_response = MagicMock()
mock_response.status_code = 401
mock_response.text = raw_response_body
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
"Unauthorized",
request=MagicMock(),
response=mock_response,
)
mock_client = AsyncMock()
mock_client.post.return_value = mock_response
with (
patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
),
patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.verbose_logger.debug"
) as mock_debug,
pytest.raises(ValueError, match="failed with status 401"),
):
await handler.exchange_token("bad-jwt", server)
logged_values = " ".join(
str(value)
for call in mock_debug.call_args_list
for value in [*call.args, *call.kwargs.values()]
)
assert raw_response_body not in logged_values
@pytest.mark.asyncio
async def test_exchange_token_missing_access_token():
"""Response without access_token raises ValueError."""
handler = TokenExchangeHandler()
server = _obo_server()
resp = MagicMock()
resp.json.return_value = {"token_type": "Bearer"}
resp.raise_for_status = MagicMock()
mock_client = AsyncMock()
mock_client.post.return_value = resp
with (
patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.get_async_httpx_client",
return_value=mock_client,
),
pytest.raises(ValueError, match="missing 'access_token'"),
):
await handler.exchange_token("jwt", server)
@pytest.mark.asyncio
async def test_exchange_token_missing_endpoint():
"""Missing token_exchange_endpoint and token_url raises ValueError."""
handler = TokenExchangeHandler()
server = _obo_server(token_exchange_endpoint=None, token_url=None)
with pytest.raises(ValueError, match="no token_exchange_endpoint or token_url"):
await handler.exchange_token("jwt", server)
@pytest.mark.asyncio
async def test_exchange_token_missing_credentials():
"""Missing client_id or client_secret raises ValueError."""
handler = TokenExchangeHandler()
server = _obo_server(client_id=None, client_secret=None)
# has_token_exchange_config will be False, so we call _do_exchange directly
with pytest.raises(ValueError, match="missing client_id or client_secret"):
await handler._do_exchange("jwt", server)
# ── resolve_mcp_auth Integration ──
@pytest.mark.asyncio
async def test_resolve_mcp_auth_with_token_exchange():
"""resolve_mcp_auth delegates to token exchange when server has OBO config and subject_token provided."""
server = _obo_server()
mock_handler = AsyncMock()
mock_handler.exchange_token.return_value = "obo-scoped-token"
with patch(
"litellm.proxy._experimental.mcp_server.auth.token_exchange.mcp_token_exchange_handler",
mock_handler,
):
result = await resolve_mcp_auth(server, subject_token="user-jwt")
assert result == "obo-scoped-token"
mock_handler.exchange_token.assert_called_once_with("user-jwt", server)
@pytest.mark.asyncio
async def test_resolve_mcp_auth_obo_without_subject_token_falls_through():
"""Without a subject_token, resolve_mcp_auth falls through to client_credentials."""
server = _obo_server(
token_url="https://auth.example.com/token",
)
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response("cc-token")
with patch(
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client",
return_value=mock_client,
):
result = await resolve_mcp_auth(server, subject_token=None)
# Falls through to client_credentials since subject_token is None
# The server has client_id/client_secret/token_url so has_client_credentials is True
assert result == "cc-token"
@pytest.mark.asyncio
async def test_resolve_mcp_auth_obo_without_subject_token_uses_cached_client_credentials():
"""The M2M fallback for OBO servers reuses the client_credentials cache."""
server = _obo_server(
server_id="srv-obo-m2m-cache",
token_url="https://auth.example.com/token",
)
mock_client = AsyncMock()
mock_client.post.return_value = _exchange_response("cached-cc-token")
with patch(
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client",
return_value=mock_client,
):
first = await resolve_mcp_auth(server, subject_token=None)
second = await resolve_mcp_auth(server, subject_token=None)
assert first == second == "cached-cc-token"
mock_client.post.assert_called_once()
@pytest.mark.asyncio
async def test_resolve_mcp_auth_header_beats_obo():
"""An explicit mcp_auth_header takes priority over OBO token exchange."""
server = _obo_server()
result = await resolve_mcp_auth(
server, mcp_auth_header="Bearer override", subject_token="user-jwt"
)
assert result == "Bearer override"
# ── Bearer Token Extraction ──
def test_extract_bearer_token_from_oauth2_headers():
"""Extracts token from oauth2_headers Authorization header."""
result = MCPServerManager._extract_bearer_token(
oauth2_headers={"Authorization": "Bearer my-jwt-token"},
raw_headers=None,
)
assert result == "my-jwt-token"
def test_extract_bearer_token_from_raw_headers():
"""Falls back to raw_headers when oauth2_headers missing."""
result = MCPServerManager._extract_bearer_token(
oauth2_headers=None,
raw_headers={"authorization": "Bearer raw-jwt"},
)
assert result == "raw-jwt"
def test_extract_bearer_token_no_bearer_prefix():
"""Returns token as-is when no Bearer prefix."""
result = MCPServerManager._extract_bearer_token(
oauth2_headers={"Authorization": "some-opaque-token"},
raw_headers=None,
)
assert result == "some-opaque-token"
def test_extract_bearer_token_none():
"""Returns None when no auth headers present."""
result = MCPServerManager._extract_bearer_token(
oauth2_headers=None,
raw_headers=None,
)
assert result is None
# ── MCPServer Properties ──
def test_has_token_exchange_config_true():
"""has_token_exchange_config is True for a fully configured OBO server."""
server = _obo_server()
assert server.has_token_exchange_config is True
def test_has_token_exchange_config_false_wrong_auth_type():
"""has_token_exchange_config is False when auth_type is not oauth2_token_exchange."""
server = _obo_server(auth_type=MCPAuth.oauth2)
assert server.has_token_exchange_config is False
def test_has_token_exchange_config_false_missing_creds():
"""has_token_exchange_config is False when client_id/client_secret missing."""
server = _obo_server(client_id=None)
assert server.has_token_exchange_config is False
def test_has_token_exchange_config_uses_token_url_fallback():
"""has_token_exchange_config is True when token_url is set instead of token_exchange_endpoint."""
server = _obo_server(
token_exchange_endpoint=None,
token_url="https://idp.example.com/token",
)
assert server.has_token_exchange_config is True
# ── Config Loading ──
@pytest.mark.asyncio
async def test_config_loading_token_exchange_fields():
"""load_servers_from_config correctly maps OBO config fields to MCPServer."""
manager = MCPServerManager()
config = {
"my_obo_server": {
"url": "https://mcp.example.com/mcp",
"transport": "http",
"auth_type": "oauth2_token_exchange",
"client_id": "my-client",
"client_secret": "my-secret",
"token_exchange_endpoint": "https://idp.example.com/oauth2/token",
"audience": "api://my-mcp",
"scopes": ["read", "write"],
"subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
}
}
await manager.load_servers_from_config(config)
servers = list(manager.config_mcp_servers.values())
assert len(servers) == 1
server = servers[0]
assert server.auth_type == MCPAuth.oauth2_token_exchange
assert server.token_exchange_endpoint == "https://idp.example.com/oauth2/token"
assert server.audience == "api://my-mcp"
assert server.subject_token_type == "urn:ietf:params:oauth:token-type:jwt"
assert server.client_id == "my-client"
assert server.client_secret == "my-secret"
assert server.scopes == ["read", "write"]
assert server.has_token_exchange_config is True
@pytest.mark.asyncio
async def test_config_loading_default_subject_token_type():
"""subject_token_type defaults to access_token when not specified in config."""
manager = MCPServerManager()
config = {
"obo_defaults": {
"url": "https://mcp.example.com/mcp",
"transport": "http",
"auth_type": "oauth2_token_exchange",
"client_id": "cid",
"client_secret": "csec",
"token_exchange_endpoint": "https://idp.example.com/token",
}
}
await manager.load_servers_from_config(config)
server = list(manager.config_mcp_servers.values())[0]
assert server.subject_token_type == "urn:ietf:params:oauth:token-type:access_token"
@pytest.mark.asyncio
async def test_database_loading_token_exchange_scopes_from_credentials():
"""DB-loaded OBO server credentials retain configured scopes."""
manager = MCPServerManager()
db_server = LiteLLM_MCPServerTable(
server_id="srv-obo-db",
server_name="obo_db_server",
url="https://mcp.example.com/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
credentials={
"client_id": "db-client",
"client_secret": "db-secret",
"token_exchange_endpoint": "https://idp.example.com/oauth2/token",
"audience": "api://db-mcp",
"scopes": ["db.read", "db.write"],
},
)
server = await manager.build_mcp_server_from_table(
db_server,
credentials_are_encrypted=False,
)
assert server.auth_type == MCPAuth.oauth2_token_exchange
assert server.client_id == "db-client"
assert server.client_secret == "db-secret"
assert server.token_exchange_endpoint == "https://idp.example.com/oauth2/token"
assert server.audience == "api://db-mcp"
assert server.scopes == ["db.read", "db.write"]

30
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py
View File

@ -549,7 +549,11 @@ class TestHookHeaderMergePriority:
captured_extra_headers: Dict[str, Any] = {}
async def fake_create_mcp_client(
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
server,
mcp_auth_header=None,
extra_headers=None,
stdio_env=None,
subject_token=None,
):
captured_extra_headers["value"] = extra_headers
mock_client = MagicMock()
@ -589,7 +593,11 @@ class TestHookHeaderMergePriority:
captured_extra_headers: Dict[str, Any] = {}
async def fake_create_mcp_client(
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
server,
mcp_auth_header=None,
extra_headers=None,
stdio_env=None,
subject_token=None,
):
captured_extra_headers["value"] = extra_headers
mock_client = MagicMock()
@ -635,7 +643,11 @@ class TestHookHeaderMergePriority:
captured_extra_headers: Dict[str, Any] = {}
async def fake_create_mcp_client(
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
server,
mcp_auth_header=None,
extra_headers=None,
stdio_env=None,
subject_token=None,
):
captured_extra_headers["value"] = extra_headers
mock_client = MagicMock()
@ -691,7 +703,11 @@ class TestHookHeaderMergePriority:
captured_extra_headers: Dict[str, Any] = {}
async def fake_create_mcp_client(
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
server,
mcp_auth_header=None,
extra_headers=None,
stdio_env=None,
subject_token=None,
):
captured_extra_headers["value"] = extra_headers
mock_client = MagicMock()
@ -739,7 +755,11 @@ class TestHookHeaderMergePriority:
captured_extra_headers: Dict[str, Any] = {}
async def fake_create_mcp_client(
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
server,
mcp_auth_header=None,
extra_headers=None,
stdio_env=None,
subject_token=None,
):
captured_extra_headers["value"] = extra_headers
mock_client = MagicMock()

2
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View File

@ -1228,6 +1228,7 @@ async def test_oauth2_headers_passed_to_mcp_client():
mcp_auth_header=None,
extra_headers=None,
stdio_env=None,
subject_token=None,
):
# Capture the arguments for verification
captured_client_args.update(
@ -1236,6 +1237,7 @@ async def test_oauth2_headers_passed_to_mcp_client():
"mcp_auth_header": mcp_auth_header,
"extra_headers": extra_headers,
"stdio_env": stdio_env,
"subject_token": subject_token,
}
)
# Return a mock client that doesn't actually connect

2
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View File

@ -450,7 +450,7 @@ class TestMCPServerManager:
captured_extra_headers = None
async def capture_create_mcp_client(
server, mcp_auth_header, extra_headers, stdio_env
server, mcp_auth_header, extra_headers, stdio_env, subject_token=None
): # pragma: no cover - helper
nonlocal captured_extra_headers
captured_extra_headers = extra_headers