scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-13 03:59:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge f524d452cb into 780c504e00

Browse source
This commit is contained in:
joshua-berri 2026-10-11 16:18:05 +00:00 • committed by GitHub
commit 4841753b4e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
36 changed files with 2021 additions and 123 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.github/workflows/test-unit.yml vendored
View file

@ -794,6 +794,10 @@ jobs:
tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py
tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py
tests/unit/proxy/_experimental/mcp_server/outbound_credentials
python-version: "3.12"
workers: 4

35
litellm/experimental_mcp_client/client.py
View file

@ -41,6 +41,7 @@ class _SessionCallbackMap(TypedDict):
sampling_callback: ReadOnly[SamplingFnT | None]
elicitation_callback: ReadOnly[ElicitationFnT | None]
logging_callback: ReadOnly[LoggingFnT | None]
response_hooks: ReadOnly[tuple[Callable[[httpx2.Response], Awaitable[None]], ...]]
from mcp.types import (
@ -96,6 +97,7 @@ from litellm.constants import (
)
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
from litellm.proxy._experimental.mcp_server.mcp_context import get_active_mcp_request_ctx
from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response
from litellm.proxy._experimental.mcp_server.result_conversion import (
age_freshness,
@ -122,6 +124,20 @@ if TYPE_CHECKING:
from litellm.proxy._experimental.mcp_server.legacy_callbacks import ElicitationCallback
def _mcp_response_hooks() -> tuple[Callable[[httpx2.Response], Awaitable[None]], ...]:
if get_active_mcp_request_ctx() is None:
return (capture_upstream_error_response,)
from litellm.proxy._experimental.mcp_server.legacy_callbacks import (
get_scope_response,
record_upstream_tool_authorization,
)
return (
partial(record_upstream_tool_authorization, pending=get_scope_response()),
capture_upstream_error_response,
)
def to_basic_auth(auth_value: str) -> str:
"""Convert auth value to Basic Auth format."""
return base64.b64encode(auth_value.encode("utf-8")).decode()
@ -497,6 +513,7 @@ class MCPClient:
"sampling_callback": self._sampling_callback,
"elicitation_callback": self._elicitation_callback,
"logging_callback": self._logging_callback,
"response_hooks": _mcp_response_hooks(),
}
def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool]:
@ -721,7 +738,9 @@ class MCPClient:
write_stream,
read_timeout_seconds=self.timeout,
message_handler=receive_message,
**session_kwargs,
sampling_callback=session_kwargs["sampling_callback"],
elicitation_callback=session_kwargs["elicitation_callback"],
logging_callback=session_kwargs["logging_callback"],
)
session: Final = await session_ctx.__aenter__()
try:
@ -931,7 +950,14 @@ class MCPClient:
verify=ssl_config,
follow_redirects=True,
event_hooks=MappingProxyType(
{"response": [capture_upstream_error_response], "request": [guard] if guard else []}
{
"response": list(
self._session_callbacks_override["response_hooks"]
if self._session_callbacks_override is not None
else _mcp_response_hooks()
),
"request": [guard] if guard else [],
}
),
)
@ -1582,6 +1608,7 @@ class PersistentMCPSession:
self._initial_callbacks: Final[_SessionCallbackMap] = client.session_callbacks()
self._current_callbacks: _SessionCallbackMap = self._initial_callbacks
delegated_callbacks: Final[_SessionCallbackMap] = {
"response_hooks": (self._record_response,),
"sampling_callback": (
self._delegate("sampling_callback")
if self._initial_callbacks["sampling_callback"] is not None
@ -1599,6 +1626,10 @@ class PersistentMCPSession:
client.route_session_callbacks(delegated_callbacks)
self._task: Final = loop.create_task(self._serve())
async def _record_response(self, response: httpx2.Response) -> None:
for hook in self._current_callbacks["response_hooks"]:
await hook(response)
@overload
def _delegate(self, name: Literal["sampling_callback"]) -> SamplingFnT: ...

6
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -1790,7 +1790,7 @@ async def store_user_oauth_credential(
payload["refresh_token"] = refresh_token
if expires_at:
payload["expires_at"] = expires_at
if scopes:
if scopes is not None:
payload["scopes"] = scopes
# Guard against silently overwriting a BYOK credential with an OAuth token.
@ -2123,9 +2123,7 @@ async def refresh_user_oauth_token(
new_refresh_token: Final[str | None] = body.get("refresh_token") or refresh_token
raw_scope: Final = body.get("scope")
scopes: list[str] | None = (raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None) or cred.get(
"scopes"
)
scopes: Final = raw_scope.split() if isinstance(raw_scope, str) else cred.get("scopes")
await store_user_oauth_credential(
prisma_client=prisma_client,

237
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -12,6 +12,8 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
import httpx
from fastapi import APIRouter, Depends, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from mcp.client.auth.utils import get_client_metadata_scopes, union_scopes
from mcp.shared.auth import OAuthMetadata
from pydantic import ConfigDict, Field, SecretStr, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
@ -110,6 +112,7 @@ from litellm.types.mcp import MCPAuth, MCPCredentials
from litellm.types.mcp_server.mcp_server_manager import MCPServer, MCPTokenEndpointAuthMethod
if TYPE_CHECKING:
from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload
from litellm.proxy._types import LiteLLM_MCPServerTable
# TTL cache for upstream OAuth metadata fetched from pass-through MCP servers.
@ -219,6 +222,8 @@ def encode_state_with_base_url(
expected_issuer: str | None = None,
authorization_response_iss_parameter_supported: bool = False,
oauth_nonce: str | None = None,
managed_scope: str | None = None,
scope_user_id: str | None = None,
) -> str:
"""
Encode the base_url, original state, and PKCE parameters using encryption.
@ -251,6 +256,7 @@ def encode_state_with_base_url(
"""
state_data: Final = {
"oauth_nonce": oauth_nonce,
**({"managed_scope": managed_scope, "scope_user_id": scope_user_id} if managed_scope is not None else {}),
"base_url": base_url,
"original_state": original_state,
"code_challenge": code_challenge,
@ -290,6 +296,18 @@ def decode_state_hash(encrypted_state: str) -> dict:
return state_data
_MANAGED_AUTH_CODE_PREFIX: Final = "llm_scope_"
class _ManagedAuthorizationCode(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
upstream_code: str = Field(min_length=1)
mcp_server_id: str = Field(min_length=1)
issuer: str | None = None
requested_scope: str
scope_user_id: str | None = None
_BRIDGE_AUTH_CODE_PREFIX: Final = "llm_bcode_"
@ -298,6 +316,8 @@ class _BridgeAuthorizationCode(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
oauth_nonce: str | None = None
requested_scope: str | None = None
expected_issuer: str | None = None
upstream_code: str = Field(min_length=1)
litellm_user_id: str = Field(min_length=1)
mcp_server_id: str = Field(min_length=1)
@ -314,6 +334,8 @@ def seal_bridge_authorization_code(
litellm_user_id: str,
mcp_server_id: str,
oauth_nonce: str | None = None,
requested_scope: str | None = None,
expected_issuer: str | None = None,
) -> str:
"""Seal the upstream authorization code and the SSO-captured litellm user into a gateway
authorization code. The DCR client only echoes this opaque value back at the token endpoint; the
@ -328,6 +350,8 @@ def seal_bridge_authorization_code(
"litellm_user_id": litellm_user_id,
"mcp_server_id": mcp_server_id,
"oauth_nonce": oauth_nonce,
"requested_scope": requested_scope,
"expected_issuer": expected_issuer,
},
sort_keys=True,
)
@ -643,6 +667,7 @@ async def _store_per_user_token_server_side(
token_response: dict[str, Any],
identity_binding_proof: str | None = None,
cimd_client_id: str | None = None,
requested_scope: str | None = None,
) -> None:
"""Persist the OAuth token server-side and warm the Redis cache.
@ -667,8 +692,8 @@ async def _store_per_user_token_server_side(
expires_in = None
refresh_token: Final[str | None] = token_response.get("refresh_token") or None
raw_scope: Final = token_response.get("scope")
scopes: Final[list | None] = raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None
raw_scope: Final[object] = token_response.get("scope", requested_scope)
scopes: Final[list[str] | None] = raw_scope.split() if isinstance(raw_scope, str) else None
try:
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Cannot store per-user OAuth token.")
@ -954,6 +979,46 @@ async def _resolve_oauth_authorization_user(
return user_id
async def _managed_authorization_scope(
request: Request, server: MCPServer, scope: str | None, authorized_user_id: str | None
) -> tuple[str | None, str | None]:
from litellm.proxy._experimental.mcp_server.db import get_user_oauth_credential
from litellm.proxy.proxy_server import prisma_client
gateway_credential: Final = authorized_user_id is None and await oauth_authorization_uses_gateway_credential(
request
)
session_user: Final = (
_session_cookie_user_id(request) if authorized_user_id is None and not gateway_credential else None
)
scope_user_id: Final = authorized_user_id or (
await authorize_oauth_credential_request(request, server.server_id)
if gateway_credential
else session_user
if session_user and await _user_can_reach_mcp_server(session_user, server.server_id)
else None
)
prior_grant: Final = (
await get_user_oauth_credential(prisma_client, scope_user_id, server.server_id)
if scope_user_id and prisma_client is not None
else None
)
merged_scope: Final = union_scopes(" ".join(prior_grant.get("scopes") or ()) if prior_grant else None, scope)
selected_scope: Final = (
get_client_metadata_scopes(
merged_scope,
None,
OAuthMetadata.model_validate(server.authorization_server_metadata)
if server.authorization_server_metadata is not None
else None,
["authorization_code", "refresh_token"],
)
if merged_scope is not None
else None
)
return selected_scope, scope_user_id
async def authorize_with_server(
request: Request,
mcp_server: MCPServer,
@ -1029,6 +1094,17 @@ async def authorize_with_server(
return subject
litellm_user_id = subject
managed_per_user: Final = resolved_server.is_gateway_managed_oauth2 and resolved_server.needs_user_oauth_token
initial_scope: Final = union_scopes(
"openid" if enforce_binding else None,
scope or (" ".join(resolved_server.scopes) if resolved_server.scopes else None),
)
selected_scope, scope_user_id = (
await _managed_authorization_scope(request, resolved_server, initial_scope, litellm_user_id)
if managed_per_user
else (initial_scope, None)
)
oauth_nonce: Final = secrets.token_urlsafe(32) if enforce_binding else None
encoded_state: Final = encode_state_with_base_url(
base_url=base_url,
@ -1038,7 +1114,11 @@ async def authorize_with_server(
code_challenge_method=code_challenge_method,
client_redirect_uri=redirect_uri,
litellm_user_id=litellm_user_id,
mcp_server_id=resolved_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None,
mcp_server_id=resolved_server.server_id
if (litellm_user_id or ephemeral_dcr_client or managed_per_user)
else None,
managed_scope=selected_scope if managed_per_user else None,
scope_user_id=scope_user_id,
dcr_client_id=ephemeral_dcr_client.client_id if ephemeral_dcr_client else None,
dcr_client_secret=ephemeral_dcr_client.client_secret if ephemeral_dcr_client else None,
dcr_token_endpoint_auth_method=ephemeral_dcr_client.token_endpoint_auth_method
@ -1057,13 +1137,8 @@ async def authorize_with_server(
}
if oauth_nonce:
params["nonce"] = oauth_nonce
if scope:
params["scope"] = scope
elif resolved_server.scopes:
params["scope"] = " ".join(resolved_server.scopes)
if enforce_binding and "openid" not in params.get("scope", "").split():
params["scope"] = f"openid {params.get('scope', '')}".strip()
if selected_scope:
params["scope"] = selected_scope
if code_challenge:
params["code_challenge"] = code_challenge
@ -1090,18 +1165,11 @@ def _token_credential_source(mcp_server: MCPServer) -> CredentialSource:
return "gateway_stored" if mcp_server.client_id or get_cimd_client_id(mcp_server) else "caller_supplied"
async def _saved_cimd_refresh_client_id(
async def _saved_refresh_credential(
server: MCPServer, user_id: str | None, refresh_token: str | None
) -> str | None:
"""Reuse only the client identity bound to this caller's presented refresh grant."""
if (
not user_id
or not refresh_token
or not server.needs_user_oauth_token
or server.auth_type != MCPAuth.oauth2
or server.client_id
or server.client_secret
):
) -> "OAuthCredentialPayload | None":
"""Reuse only metadata bound to this caller's presented refresh grant."""
if not user_id or not refresh_token or not server.needs_user_oauth_token or server.auth_type != MCPAuth.oauth2:
return None
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.db import get_user_oauth_credential
@ -1117,7 +1185,33 @@ async def _saved_cimd_refresh_client_id(
stored_refresh: Final = credential.get("refresh_token")
if not stored_refresh or not secrets.compare_digest(stored_refresh.encode(), refresh_token.encode()):
return None
return credential.get("cimd_client_id")
return credential
def _open_managed_authorization_code(
code: str | None, server: MCPServer, grant_type: str, user_id: str | None
) -> _ManagedAuthorizationCode | None:
if not isinstance(code, str) or not code.startswith(_MANAGED_AUTH_CODE_PREFIX):
return None
decoded_code: Final = decrypt_value_helper(
code[len(_MANAGED_AUTH_CODE_PREFIX) :], "managed_authorization_code", return_original_value=False
)
if not isinstance(decoded_code, str):
raise HTTPException(status_code=400, detail="Invalid authorization code")
try:
managed_code: Final = _ManagedAuthorizationCode.model_validate_json(decoded_code)
except ValidationError as exc:
raise HTTPException(status_code=400, detail="Invalid authorization code") from exc
if (
grant_type != "authorization_code"
or not server.is_gateway_managed_oauth2
or not server.needs_user_oauth_token
or managed_code.mcp_server_id != server.server_id
or managed_code.issuer != server.issuer
or (managed_code.scope_user_id is not None and managed_code.scope_user_id != user_id)
):
raise HTTPException(status_code=400, detail="Authorization code does not match this OAuth flow")
return managed_code
async def exchange_token_with_server(
@ -1159,10 +1253,18 @@ async def exchange_token_with_server(
if resolved_server.needs_user_oauth_token or resolved_server.oauth_identity_binding is not None
else None
)
cimd_client_id: Final = (
await _saved_cimd_refresh_client_id(resolved_server, request_user_id, refresh_token)
managed_code: Final = _open_managed_authorization_code(code, resolved_server, grant_type, request_user_id)
authorization_code: Final = managed_code.upstream_code if managed_code else code
refresh_credential: Final = (
await _saved_refresh_credential(resolved_server, request_user_id, refresh_token)
if grant_type == "refresh_token"
else None
)
cimd_client_id: Final = (
refresh_credential.get("cimd_client_id")
if refresh_credential and not resolved_server.client_id and not resolved_server.client_secret
else None
) or get_cimd_client_id(resolved_server)
# The id, secret, and token-endpoint auth method must come from the same source. When the
@ -1234,7 +1336,7 @@ async def exchange_token_with_server(
)
else:
refresh_ownership = None # rebind-ok: grant-specific branches assign one ownership value
if not code:
if not authorization_code:
raise HTTPException(
status_code=400,
detail="code is required for authorization_code grant",
@ -1244,14 +1346,18 @@ async def exchange_token_with_server(
# below uses the upstream code, and the mint binds the envelope to the recovered user. Bind the
# sealed server to this request so a code minted for one bridge server cannot be spent at another.
# A raw upstream code (scripted path) opens to None and the code is used as-is.
bridge_identity = open_bridge_authorization_code(code)
if bridge_identity is not None:
if bridge_identity.mcp_server_id != resolved_server.server_id:
raise HTTPException(
status_code=400,
detail="Authorization code was issued for a different MCP server",
)
code = bridge_identity.upstream_code
bridge_identity = open_bridge_authorization_code(authorization_code)
if bridge_identity is not None and (
bridge_identity.mcp_server_id != resolved_server.server_id
or (
bridge_identity.requested_scope is not None
and bridge_identity.expected_issuer != resolved_server.issuer
)
):
raise HTTPException(
status_code=400,
detail="Authorization code was issued for a different MCP server",
)
binding: Final = resolved_server.oauth_identity_binding
if binding is not None and binding.mode == "enforce":
if bridge_identity is None or not bridge_identity.oauth_nonce:
@ -1273,7 +1379,7 @@ async def exchange_token_with_server(
resolved_redirect_uri: Final = redirect_uri if bridge_token_relay else f"{proxy_base_url}/callback"
token_data = {
"grant_type": "authorization_code",
"code": code,
"code": bridge_identity.upstream_code if bridge_identity else authorization_code,
"redirect_uri": resolved_redirect_uri,
**token_request.body,
}
@ -1353,6 +1459,16 @@ async def exchange_token_with_server(
else None
)
requested_scope: Final = (
managed_code.requested_scope
if managed_code is not None
else bridge_identity.requested_scope
if bridge_identity is not None
else scope or " ".join(refresh_credential.get("scopes") or ())
if refresh_credential
else None
)
# Store server-side when the server is configured for per-user OAuth and
# the calling client has provided a valid LiteLLM identity.
# Errors are non-fatal: the token is still returned to the client.
@ -1379,6 +1495,7 @@ async def exchange_token_with_server(
user_id=user_id,
token_response=token_response,
identity_binding_proof=binding_proof,
requested_scope=requested_scope,
**({"cimd_client_id": cimd_client_id} if cimd_client_id is not None else {}),
)
else:
@ -1431,8 +1548,10 @@ async def exchange_token_with_server(
result["expires_in"] = token_response["expires_in"]
if token_response.get("refresh_token"):
result["refresh_token"] = token_response["refresh_token"]
if token_response.get("scope"):
result["scope"] = token_response["scope"]
token_fields: Final = TypeAdapter[dict[str, object]](dict[str, object]).validate_python(token_response)
effective_scope: Final = token_fields.get("scope", requested_scope)
if effective_scope is not None:
result["scope"] = effective_scope
# RFC 6749 §5.1: token responses must not be cached.
return JSONResponse(result, headers=TOKEN_NO_CACHE_HEADERS)
@ -2116,6 +2235,7 @@ async def authorize_mcp_session(
code_challenge_method: str | None = None,
response_type: str | None = None,
resource: str | None = None,
scope: str | None = None,
) -> Response:
return aggregate_authorize(
request=request,
@ -2127,6 +2247,7 @@ async def authorize_mcp_session(
response_type=response_type,
session_user_id=_session_cookie_user_id(request),
resource=resource,
scope=scope,
)
@ -2169,6 +2290,7 @@ async def authorize(
response_type=response_type,
session_user_id=_session_cookie_user_id(request),
resource=resource,
scope=scope,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
@ -2281,7 +2403,9 @@ async def token_endpoint(
)
async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState:
async def _vendor_credential_state(
user_id: str, server_id: str, requested_scopes: tuple[str, ...] = ()
) -> VendorCredentialState:
"""Whether the gateway itself can see a live vendor credential for this user and server.
The one reading of "authorized" the connect page displays and the finish step enforces, so
@ -2299,7 +2423,10 @@ async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCreden
credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id)
except Exception: # noqa: BLE001 # a credential-read fault must fail the scoped grant closed
return "unavailable"
return "absent" if oauth_grant_state(credential) == "absent" else "present"
if oauth_grant_state(credential) == "absent":
return "absent"
granted_scopes: Final = credential.get("scopes") if credential else None
return "present" if set(requested_scopes).issubset(granted_scopes or ()) else "absent"
@router.get("/authorize/flow")
@ -2542,6 +2669,11 @@ async def callback(
mcp_server_id: Final = state_data.get("mcp_server_id")
dcr_client_id: Final = state_data.get("dcr_client_id")
dcr_client_secret: Final = state_data.get("dcr_client_secret")
managed_scope: Final = issuer_state.get("managed_scope")
expected_issuer: Final = TypeAdapter[str | None](str | None).validate_python(
issuer_state.get("expected_issuer")
)
scope_owner: Final = TypeAdapter[str | None](str | None).validate_python(issuer_state.get("scope_user_id"))
forwarded_code = code
if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id:
forwarded_code = seal_bridge_authorization_code(
@ -2549,6 +2681,8 @@ async def callback(
litellm_user_id=litellm_user_id,
mcp_server_id=mcp_server_id,
oauth_nonce=state_data.get("oauth_nonce"),
requested_scope=managed_scope if isinstance(managed_scope, str) else None,
expected_issuer=expected_issuer,
)
elif isinstance(dcr_client_id, str) and dcr_client_id and isinstance(mcp_server_id, str) and mcp_server_id:
forwarded_code = seal_passthrough_authorization_code(
@ -2561,7 +2695,23 @@ async def callback(
),
)
params = {"code": forwarded_code, "state": original_state}
managed_code: Final = (
_ManagedAuthorizationCode(
upstream_code=forwarded_code,
mcp_server_id=mcp_server_id,
issuer=expected_issuer,
requested_scope=managed_scope,
scope_user_id=scope_owner,
)
if isinstance(managed_scope, str) and isinstance(mcp_server_id, str) and not litellm_user_id
else None
)
returned_code: Final = (
_MANAGED_AUTH_CODE_PREFIX + encrypt_value_helper(managed_code.model_dump_json())
if managed_code is not None
else forwarded_code
)
params = {"code": returned_code, "state": original_state}
complete_returned_url = _append_query_params(redirect_uri, params)
response = RedirectResponse(url=complete_returned_url, status_code=302)
_clear_oauth_state_cookie(response, request, state)
@ -2807,11 +2957,16 @@ async def _build_oauth_protected_resource_response(
if obo_response is not None:
return obo_response
resource_scopes: Final = [
value
for value in (mcp_server.scopes or () if mcp_server else ())
if value != "offline_access" or not mcp_server or not mcp_server.is_gateway_managed_oauth2
]
if mcp_server is not None and mcp_server.advertises_gateway_authorization_server:
return {
"authorization_servers": [f"{request_base_url}/mcp"],
"resource": resource_url,
"scopes_supported": (mcp_server.scopes if mcp_server.scopes else []),
"scopes_supported": resource_scopes,
}
if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange:
@ -2822,7 +2977,7 @@ async def _build_oauth_protected_resource_response(
(f"{request_base_url}/{mcp_server_name}" if mcp_server_name else f"{request_base_url}")
],
"resource": resource_url,
"scopes_supported": (mcp_server.scopes if mcp_server and mcp_server.scopes else []),
"scopes_supported": resource_scopes,
}

32
litellm/proxy/_experimental/mcp_server/exceptions.py
View file

@ -4,6 +4,8 @@ from typing import Final
from fastapi import HTTPException
from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope
class MCPServerURLCredentialsError(HTTPException):
"""A fixed, sanitized URL-credential migration error safe for operator previews."""
@ -37,10 +39,14 @@ class MCPUpstreamAuthError(Exception):
status_code: int,
www_authenticate: str | None,
server_name: str,
required_scope: str | None = None,
server_id: str | None = None,
) -> None:
self.status_code = status_code
self.www_authenticate = www_authenticate
self.server_name = server_name
self.required_scope = required_scope
self.server_id = server_id
super().__init__(f"Upstream MCP server {server_name!r} returned {status_code}")
def to_http_exception(
@ -73,19 +79,39 @@ class MCPUpstreamAuthError(Exception):
the client originally targeted, matching the path-aware behaviour of
``get_passthrough_resource_metadata_url`` in ``oauth_utils.py``.
"""
aggregate_step_up: Final = (
self.required_scope is not None
and request_path is not None
and (request_path.rstrip("/") == "/mcp" or request_path.startswith("/toolset/"))
)
requested_scope: Final = (
" ".join(
encode_upstream_scope(self.server_id or self.server_name, item) for item in self.required_scope.split()
)
if aggregate_step_up and self.required_scope
else self.required_scope
)
challenge: str | None = self.www_authenticate
if challenge is None and self.status_code == 401 and base_url:
if challenge is None and (self.status_code == 401 or self.required_scope is not None) and base_url:
prefix: Final = base_url.rstrip("/")
if request_path and request_path.startswith(f"/{self.server_name}/mcp"):
if aggregate_step_up:
resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource"
elif request_path and request_path.startswith(f"/{self.server_name}/mcp"):
resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/{self.server_name}/mcp"
else:
resource_metadata_url = f"{prefix}/.well-known/oauth-protected-resource/mcp/{self.server_name}"
challenge = f'Bearer resource_metadata="{resource_metadata_url}"'
scoped_challenge: Final = (
(f'{challenge}, error="insufficient_scope"' if challenge else 'Bearer error="insufficient_scope"')
+ (f', scope="{requested_scope}"' if requested_scope else "")
if self.required_scope is not None
else challenge
)
detail: Final = "Forbidden" if self.status_code == 403 else "Unauthorized"
return HTTPException(
status_code=self.status_code,
detail=detail,
headers={"www-authenticate": challenge} if challenge else None,
headers={"www-authenticate": scoped_challenge} if scoped_challenge else None,
)

27
litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py
View file

@ -10,13 +10,15 @@ becomes an outcome, never a second failure.
from __future__ import annotations
import re
from collections.abc import Iterator
from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias
from urllib.request import parse_http_list, parse_keqv_list
import httpx
import httpx2
from mcp.types import Tool as MCPTool
from pydantic import ConfigDict
from pydantic import ConfigDict, TypeAdapter
from typing_extensions import assert_never
from litellm.proxy._experimental.mcp_server.exceptions import (
@ -98,6 +100,29 @@ def upstream_auth_challenge(exc: BaseException) -> tuple[int, str | None] | None
return None
def upstream_insufficient_scope(exc: BaseException) -> str | None:
auth: Final = upstream_auth_challenge(exc)
if auth is None or auth[0] != 403 or not auth[1]:
return None
scheme, _, parameters = auth[1].partition(" ")
if scheme.lower() != "bearer":
return None
fields: Final = parse_http_list(parameters)
if any("=" not in field or not field.split("=", 1)[1] for field in fields):
return None
keys: Final = tuple(field.split("=", 1)[0].strip().lower() for field in fields)
if len(set(keys)) != len(keys) or any(re.fullmatch(r"[!#$%&'*+.^_`|~0-9a-z-]+", key) is None for key in keys):
return None
parsed: Final = TypeAdapter(dict[str, str]).validate_python(parse_keqv_list(fields))
values: Final = {key: value for key, value in zip(keys, parsed.values())}
if values.get("error") != "insufficient_scope":
return None
scope: Final = values.get("scope", "")
if scope and re.fullmatch(r"[!#-\[\]-~]+(?: [!#-\[\]-~]+)*", scope) is None:
return None
return scope
def raise_classified_list_failure(
exc: BaseException,
server_name: str,

116
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -59,8 +59,11 @@ from litellm.caching.caching import DualCache
from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager
from litellm.proxy._experimental.mcp_server.oauth_utils import (
TOKEN_NO_CACHE_HEADERS,
UPSTREAM_SCOPE_PREFIX,
UpstreamScope,
canonical_resource_uri,
canonicalize_url_identity,
decode_upstream_scope,
get_request_base_url,
is_loopback_redirect_host,
validate_redirect_uri_shape,
@ -250,7 +253,9 @@ class ConsentTeam(LiteLLMBaseModel):
class LookupVendorCredential(Protocol):
"""Injected read of a user's vendor credential for one server."""
def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ...
def __call__(
self, user_id: str, server_id: str, requested_scopes: tuple[str, ...] = (), /
) -> Awaitable[VendorCredentialState]: ...
class LookupServerReachability(Protocol):
@ -273,7 +278,9 @@ async def _refuse_subject_token(subject_token: str, request: Request) -> Subject
)
async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState:
async def _unavailable_vendor_credential(
user_id: str, server_id: str, requested_scopes: tuple[str, ...] = ()
) -> VendorCredentialState:
return "unavailable"
@ -310,6 +317,8 @@ class _ConnectFlow(LiteLLMBaseModel):
exp: int
resource_server_id: str | None = None
audience: SessionAudience | None = None
requested_scopes: tuple[str, ...] = ()
upstream_scopes: tuple[UpstreamScope, ...] = ()
class _GatewayAuthCode(LiteLLMBaseModel):
@ -555,6 +564,7 @@ def aggregate_authorize(
response_type: str | None,
session_user_id: str | None,
resource: str | None = None,
scope: str | None = None,
) -> Response:
"""The aggregate authorize verb: validate the client, require S256 PKCE, interpose
LiteLLM sign-in, and hand the browser to the connect page with the flow sealed into a
@ -580,6 +590,13 @@ def aggregate_authorize(
if session_user_id is None:
return _login_redirect(base_url, request)
scoped_server: Final = resolve_scoped_resource_server(request, resource)
scope_tokens: Final = tuple(token for token in (scope or "").split() if token.startswith(UPSTREAM_SCOPE_PREFIX))
decoded_scopes: Final = tuple(decode_upstream_scope(token) for token in scope_tokens)
if any(item is None for item in decoded_scopes):
return _oauth_error(400, "invalid_scope", "invalid upstream permission scope")
upstream_scopes: Final = tuple(item for item in decoded_scopes if item is not None)
if scoped_server is not None and any(item.server_id != scoped_server.server_id for item in upstream_scopes):
return _oauth_error(400, "invalid_scope", "upstream permission does not match the requested resource")
handle: Final = secrets.token_urlsafe(24)
flow: Final = _new_connect_flow(
session_user_id=session_user_id,
@ -589,6 +606,8 @@ def aggregate_authorize(
code_challenge=code_challenge or "",
resource_server_id=scoped_server.server_id if scoped_server is not None else None,
audience=None,
requested_scopes=tuple((scope or "").split()) if scoped_server is not None and not upstream_scopes else (),
upstream_scopes=upstream_scopes,
)
connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),))
response: Final = RedirectResponse(connect_url, status_code=303)
@ -748,6 +767,8 @@ def _new_connect_flow(
code_challenge: str,
resource_server_id: str | None,
audience: SessionAudience | None,
requested_scopes: tuple[str, ...] = (),
upstream_scopes: tuple[UpstreamScope, ...] = (),
) -> _ConnectFlow:
now: Final = datetime.now(timezone.utc)
return _ConnectFlow(
@ -760,6 +781,8 @@ def _new_connect_flow(
exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS,
resource_server_id=resource_server_id,
audience=audience,
requested_scopes=requested_scopes,
upstream_scopes=upstream_scopes,
)
@ -815,20 +838,20 @@ def _open_flow_for(
@catalog_operation(global_manager)
async def _flow_target(
flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability
user_id: str, server_id: str | None, lookup_server_reachability: LookupServerReachability
) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]:
if flow.resource_server_id is None:
if server_id is None:
return "unscoped", None
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle
MCPServerManager,
global_mcp_server_manager,
)
server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id)
server: Final = global_mcp_server_manager.get_mcp_server_by_id(server_id)
if (
server is None
or not (server.is_gateway_managed_oauth2 or server.advertises_gateway_authorization_server)
or not await lookup_server_reachability(flow.user_id, server.server_id)
or not await lookup_server_reachability(user_id, server.server_id)
):
return "stale", None
state: Final = (
@ -847,6 +870,7 @@ class ConnectFlowDescription(TypedDict):
server_id: ReadOnly[str | None]
server_name: ReadOnly[str | None]
connected: ReadOnly[bool | None]
requested_scopes: NotRequired[ReadOnly[tuple[str, ...]]]
async def _describe_opened_flow(
@ -854,19 +878,11 @@ async def _describe_opened_flow(
lookup_vendor_credential: LookupVendorCredential,
lookup_server_reachability: LookupServerReachability,
) -> ConnectFlowDescription | Response:
state, server = await _flow_target(flow, lookup_server_reachability)
if flow.upstream_scopes:
return await _describe_upstream_scope_flow(flow, lookup_vendor_credential, lookup_server_reachability)
state, server = await _flow_target(flow.user_id, flow.resource_server_id, lookup_server_reachability)
if state == "interactive" and server is not None:
credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id)
if credential == "unavailable":
return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION)
interactive_description: Final[ConnectFlowDescription] = {
"state": state,
"client_origin": _origin_only(flow.redirect_uri),
"server_id": server.server_id,
"server_name": server.server_name or server.alias or server.name,
"connected": credential == "present",
}
return interactive_description
return await _describe_interactive_target(flow, server, flow.requested_scopes, lookup_vendor_credential)
described: Final[ConnectFlowDescription] = {
"state": state,
"client_origin": _origin_only(flow.redirect_uri),
@ -877,6 +893,70 @@ async def _describe_opened_flow(
return described
async def _describe_interactive_target(
flow: _ConnectFlow,
server: MCPServer,
scopes: tuple[str, ...],
lookup_vendor_credential: LookupVendorCredential,
) -> ConnectFlowDescription | Response:
credential: Final = (
await lookup_vendor_credential(flow.user_id, server.server_id, scopes)
if scopes
else await lookup_vendor_credential(flow.user_id, server.server_id)
)
if credential == "unavailable":
return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION)
described: Final[ConnectFlowDescription] = {
"state": "interactive",
"client_origin": _origin_only(flow.redirect_uri),
"server_id": server.server_id,
"server_name": server.server_name or server.alias or server.name,
"connected": credential == "present",
}
return {**described, "requested_scopes": scopes} if scopes else described
async def _describe_upstream_scope_flow(
flow: _ConnectFlow,
lookup_vendor_credential: LookupVendorCredential,
lookup_server_reachability: LookupServerReachability,
) -> ConnectFlowDescription | Response:
server_ids: Final = tuple(dict.fromkeys(item.server_id for item in flow.upstream_scopes))
targets: Final = tuple(
[await _flow_target(flow.user_id, server_id, lookup_server_reachability) for server_id in server_ids]
)
if any(state != "interactive" or server is None for state, server in targets):
return _oauth_error(400, "invalid_scope", "an upstream permission target is unavailable")
descriptions: Final = tuple(
[
await _describe_interactive_target(
flow,
server,
tuple(
dict.fromkeys(
item.permission for item in flow.upstream_scopes if item.server_id == server.server_id
)
),
lookup_vendor_credential,
)
for _, server in targets
if server is not None
]
)
for description in descriptions:
if isinstance(description, Response) or not description["connected"]:
return description
if flow.resource_server_id is not None:
return descriptions[0]
return {
"state": "unscoped",
"client_origin": _origin_only(flow.redirect_uri),
"server_id": None,
"server_name": None,
"connected": None,
}
async def describe_connect_flow(
request: Request,
flow_handle: str,

22
litellm/proxy/_experimental/mcp_server/legacy_callbacks.py
View file

@ -1,6 +1,7 @@
from collections.abc import Mapping
from typing import Final, Protocol
import httpx2
from mcp.client.session import ClientRequestContext
from mcp.types import (
CreateMessageRequestParams,
@ -14,6 +15,11 @@ from mcp.types import (
from litellm.constants import MCP_CLIENT_TIMEOUT
from litellm.proxy._experimental.mcp_server.contracts import OperationContext
from litellm.proxy._experimental.mcp_server.mcp_context import get_active_mcp_request_ctx
from litellm.proxy._experimental.mcp_server.scope_challenge import (
OAuthScopeResponse,
is_tool_call_request,
scope_response_from_request,
)
from litellm.proxy._types import UserAPIKeyAuth
@ -100,3 +106,19 @@ def create_elicitation_callback(timeout: float | None = None) -> ElicitationCall
)
return callback
def get_scope_response() -> OAuthScopeResponse | None:
context: Final = get_active_mcp_request_ctx()
return scope_response_from_request(context.request) if context is not None else None
async def record_upstream_tool_authorization(response: httpx2.Response, *, pending: OAuthScopeResponse | None) -> None:
if pending is None or not pending.pending or not response.is_success or response.request.method != "POST":
return
try:
is_tool_call: Final = is_tool_call_request(response.request.content)
except httpx2.RequestNotRead:
return
if is_tool_call:
pending.allow()

63
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -38,6 +38,7 @@ import httpx2
from fastapi import HTTPException
from httpx import HTTPStatusError
from mcp import ReadResourceResult, Resource
from mcp.shared.auth import OAuthMetadata
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import (
CallToolResult,
@ -57,7 +58,7 @@ from mcp.types import (
ResourceTemplate,
)
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl, Field, SecretStr, TypeAdapter
from pydantic import AnyUrl, Field, SecretStr, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, assert_never
import litellm
@ -103,9 +104,10 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
ServerListFault,
raise_classified_list_failure,
upstream_auth_challenge,
upstream_insufficient_scope,
)
from litellm.proxy._experimental.mcp_server.interactions import bind_target
from litellm.proxy._experimental.mcp_server.legacy_callbacks import has_active_mcp_request_context
from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response, has_active_mcp_request_context
from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure, record_auth_resolution
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
MCPPerUserTokenCache,
@ -895,6 +897,10 @@ def carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serve
new_server.authorization_response_iss_parameter_supported = ( # rebind-ok: publish on the existing rebuild object
previous_server.authorization_response_iss_parameter_supported
)
if may_carry and new_server.issuer == previous_server.issuer and new_server.authorization_server_metadata is None:
new_server.authorization_server_metadata = ( # rebind-ok: publish capability on the existing rebuild object
previous_server.authorization_server_metadata
)
if new_server.authorization_url is None and previous_server.authorization_url:
new_server.authorization_url = previous_server.authorization_url
if may_carry and new_server.token_url is None and previous_server.token_url:
@ -930,7 +936,12 @@ def _restrict_discovery_to_corroborated_authorization_server(
return metadata
if _endpoints_corroborate_authorization_url(metadata.authorization_url, manual_authorization_url):
return metadata
if not metadata.token_url and not metadata.registration_url and not metadata.client_id_metadata_document_supported:
if (
not metadata.token_url
and not metadata.registration_url
and not metadata.client_id_metadata_document_supported
and metadata.authorization_server_metadata is None
):
return metadata
bridge_note: Final = (
" The discovered registration_url is rejected with it, so this dcr_bridge server stays on the"
@ -949,7 +960,12 @@ def _restrict_discovery_to_corroborated_authorization_server(
bridge_note,
)
return metadata.model_copy(
update={"token_url": None, "registration_url": None, "client_id_metadata_document_supported": False}
update={
"token_url": None,
"registration_url": None,
"client_id_metadata_document_supported": False,
"authorization_server_metadata": None,
}
)
@ -2086,6 +2102,7 @@ class MCPServerManager:
return server.model_copy(
update={
"scopes": server.scopes or metadata.scopes,
"authorization_server_metadata": metadata.authorization_server_metadata,
"issuer": server.issuer or discovered_issuer,
"client_id_metadata_document_supported": metadata.client_id_metadata_document_supported,
"authorization_response_iss_parameter_supported": (
@ -2678,6 +2695,9 @@ class MCPServerManager:
scopes=resolved_scopes,
configured_scopes=tuple(configured_scopes) if configured_scopes else None,
issuer=effective_issuer,
authorization_server_metadata=(
gated_oauth_metadata.authorization_server_metadata if gated_oauth_metadata else None
),
client_id_metadata_document_supported=(
gated_oauth_metadata.client_id_metadata_document_supported if gated_oauth_metadata else None
),
@ -3259,6 +3279,9 @@ class MCPServerManager:
scopes=resolved_scopes,
configured_scopes=configured_scopes,
issuer=effective_issuer,
authorization_server_metadata=(
gated_oauth_metadata.authorization_server_metadata if gated_oauth_metadata else None
),
client_id_metadata_document_supported=(
gated_oauth_metadata.client_id_metadata_document_supported if gated_oauth_metadata else None
),
@ -5383,7 +5406,16 @@ class MCPServerManager:
continue
scopes = self._extract_scopes(data.get("scopes_supported"))
try:
authorization_metadata = OAuthMetadata.model_validate(data)
except ValidationError:
authorization_metadata = None
metadata = MCPOAuthMetadata(
authorization_server_metadata=(
TypeAdapter(dict[str, object]).validate_python(authorization_metadata.model_dump(mode="json"))
if authorization_metadata is not None
else None
),
scopes=scopes,
authorization_url=data.get("authorization_endpoint"),
token_url=data.get("token_endpoint"),
@ -6491,11 +6523,12 @@ class MCPServerManager:
# oauth2 + delegate_auth_to_upstream (is_oauth_passthrough) is being removed, so it is not
# added here even though the list path still relays for it.
relays_upstream_auth: Final = mcp_server.is_client_forwarded_token
managed_per_user: Final = mcp_server.is_gateway_managed_oauth2 and mcp_server.needs_user_oauth_token
server_label: Final = mcp_server.name or mcp_server.server_name or mcp_server.alias or ""
async def _call_tool_via_client(client, params):
async with self._outbound_limit(mcp_server, persistent_session):
if not relays_upstream_auth:
if not relays_upstream_auth and not managed_per_user:
return await client.call_tool(
params,
host_progress_callback=host_progress_callback,
@ -6520,8 +6553,17 @@ class MCPServerManager:
persistent_session=persistent_session,
)
except Exception as e:
required_scope: Final = upstream_insufficient_scope(e) if managed_per_user else None
if required_scope is not None:
raise MCPUpstreamAuthError(
status_code=403,
www_authenticate=None,
server_name=mcp_server.alias or server_label,
required_scope=required_scope,
server_id=mcp_server.server_id,
) from e
auth_info: Final = _extract_upstream_auth_failure(e)
if auth_info is None or auth_info[0] != 401:
if not relays_upstream_auth or auth_info is None or auth_info[0] != 401:
# A genuine (non-auth or 403-forbidden) upstream/transport failure.
# raise_on_error demoted the client-layer log to debug, so surface it here at
# warning level to keep the outage visible; the caller still gets the graceful
@ -6866,6 +6908,15 @@ class MCPServerManager:
)
tasks.append(during_hook_task)
scope_response: Final = get_scope_response()
if scope_response is not None and not (
mcp_server.is_gateway_managed_oauth2
and mcp_server.needs_user_oauth_token
and not mcp_server.spec_path
and mcp_server.transport in (MCPTransport.http, MCPTransport.sse)
):
scope_response.allow()
caller_oauth2_headers: Final = oauth2_headers
oauth2_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, oauth2_headers, user_api_key_auth)

31
litellm/proxy/_experimental/mcp_server/oauth_utils.py
View file

@ -2,6 +2,8 @@
(BYOK + discoverable / pass-through OAuth proxy)."""
import os
import re
from dataclasses import dataclass
from ipaddress import ip_address
from typing import TYPE_CHECKING, Final, NoReturn
from urllib.parse import ParseResult, urlparse, urlsplit, urlunparse, urlunsplit
@ -819,3 +821,32 @@ def build_upstream_oauth2_token_request(
if not resource:
return client_auth
return TokenEndpointClientAuth(headers=client_auth.headers, body={**client_auth.body, "resource": resource})
UPSTREAM_SCOPE_PREFIX: Final = "litellm:mcp_scope:"
@dataclass(frozen=True, slots=True)
class UpstreamScope:
server_id: str
permission: str
def encode_upstream_scope(server_id: str, permission: str) -> str:
return f"{UPSTREAM_SCOPE_PREFIX}{server_id.encode().hex()}:{permission.encode().hex()}"
def decode_upstream_scope(token: str) -> UpstreamScope | None:
encoded: Final = token.removeprefix(UPSTREAM_SCOPE_PREFIX).split(":")
if not token.startswith(UPSTREAM_SCOPE_PREFIX) or len(encoded) != 2:
return None
try:
server_id: Final = bytes.fromhex(encoded[0]).decode()
permission: Final = bytes.fromhex(encoded[1]).decode()
except (ValueError, UnicodeError):
return None
if not server_id or re.fullmatch(r"[!#-\[\]-~]+", permission) is None:
return None
if encode_upstream_scope(server_id, permission) != token:
return None
return UpstreamScope(server_id, permission)

12
litellm/proxy/_experimental/mcp_server/operations.py
View file

@ -94,6 +94,7 @@ from litellm.proxy._experimental.mcp_server.interactions import (
seal_continuation,
target_digest,
)
from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: F401 # legacy module exports
MCPServerManager,
_caller_authorization_fans_out, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@ -3009,11 +3010,12 @@ async def _execute_mcp_server_tool_call(
is_error=True,
)
except MCPUpstreamAuthError as e:
# The MCP session manager serializes handler exceptions as JSON-RPC errors, so a
# mid-session tool call cannot emit a raw 401 + WWW-Authenticate the way the REST
# call path and the connect-time preemptive check do. Return an explicit isError
# naming the upstream status (at info level, not a traceback) so the client still
# learns it must re-authenticate upstream and expected pass-through 401s don't spam.
scope_response: Final = get_scope_response()
if scope_response is not None:
scope_response.deny(e)
# The request-owned response adapter relays managed scope challenges before
# Streamable HTTP headers commit. Other auth failures retain the existing
# explicit tool error, including failures after streaming has begun.
verbose_logger.info("Upstream auth failure calling MCP tool: HTTP %s", e.status_code)
return CallToolResult(
content=[

7
litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py
View file

@ -65,7 +65,7 @@ def _parse_expires_in(raw: object) -> int | None:
def _parse_scopes(raw: object) -> tuple[str, ...] | None:
return tuple(raw.split()) if isinstance(raw, str) and raw else None
return tuple(raw.split()) if isinstance(raw, str) else None
class AuthorizationCodeRefresher:
@ -156,7 +156,8 @@ class AuthorizationCodeRefresher:
rotated: Final = body.get("refresh_token")
new_refresh: Final = rotated if isinstance(rotated, str) and rotated else token.refresh_token
expires_in: Final = _parse_expires_in(body.get("expires_in"))
scopes: Final = _parse_scopes(body.get("scope")) or token.scopes
granted_scopes: Final = _parse_scopes(body.get("scope"))
scopes: Final = granted_scopes if granted_scopes is not None else token.scopes
await self._persist(
user_id,
@ -164,7 +165,7 @@ class AuthorizationCodeRefresher:
access_token,
new_refresh,
expires_in,
scopes or None,
scopes,
**({"identity_binding_proof": binding_proof} if binding_proof is not None else {}),
**({"cimd_client_id": cimd_client_id} if cimd_client_id is not None else {}),
)

2
litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py
View file

@ -88,7 +88,7 @@ async def _persist_credential(
access_token=access_token,
refresh_token=refresh_token,
expires_in=expires_in,
scopes=list(scopes) if scopes else None,
scopes=list(scopes) if scopes is not None else None,
skip_byok_guard=True,
identity_binding_proof=identity_binding_proof,
**({"cimd_client_id": cimd_client_id} if cimd_client_id is not None else {}),

93
litellm/proxy/_experimental/mcp_server/scope_challenge.py Normal file
View file

@ -0,0 +1,93 @@
"""Keep upstream scope challenges ahead of Streamable HTTP response headers."""
from __future__ import annotations
import asyncio
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter, ValidationError
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.types import Message, Send
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
if TYPE_CHECKING:
from mcp.server.context import CallNext, HandlerResult, ServerRequestContext
SCOPE_RESPONSE_KEY: Final = "litellm.mcp.scope_response"
_REQUEST: Final = TypeAdapter[dict[str, object]](dict[str, object])
def is_tool_call_request(body: bytes) -> bool:
try:
return _REQUEST.validate_json(body).get("method") == "tools/call"
except ValidationError:
return False
class OAuthScopeResponse:
def __init__(self, send: Send, base_url: str, request_path: str | None) -> None:
self._send = send
self._base_url = base_url
self._request_path = request_path
self._ready = asyncio.Event()
self._lock = asyncio.Lock()
self._start: Message | None = None
self._challenge: MCPUpstreamAuthError | None = None
self._refused = False
@property
def pending(self) -> bool:
return not self._ready.is_set()
def allow(self) -> None:
self._ready.set()
def deny(self, error: MCPUpstreamAuthError) -> None:
if self.pending and error.status_code == 403 and error.required_scope is not None:
self._challenge = error
self._ready.set()
async def __call__(self, message: Message) -> None:
if message["type"] == "http.response.start":
self._start = message
if message["status"] != 200:
self.allow()
return
await self._ready.wait()
async with self._lock:
if self._refused:
return
if self._start is not None:
start: Final = self._start
self._start = None
if self._challenge is not None:
error: Final = self._challenge.to_http_exception(self._base_url, self._request_path)
response: Final = JSONResponse(
status_code=error.status_code, content={"detail": error.detail}, headers=error.headers
)
self._refused = True
await self._send(
{"type": "http.response.start", "status": response.status_code, "headers": response.raw_headers}
)
await self._send({"type": "http.response.body", "body": bytes(response.body)})
return
await self._send(start)
await self._send(message)
def scope_response_from_request(request: object) -> OAuthScopeResponse | None:
if not isinstance(request, Request):
return None
response: Final = request.scope.get(SCOPE_RESPONSE_KEY)
return response if isinstance(response, OAuthScopeResponse) else None
async def finish_scope_response(context: ServerRequestContext[object, object], call_next: CallNext) -> HandlerResult:
response: Final = scope_response_from_request(context.request)
try:
return await call_next(context)
finally:
if response is not None:
response.allow()

22
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -72,6 +72,12 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
redact_mcp_resource_url,
well_known_root_suffix,
)
from litellm.proxy._experimental.mcp_server.scope_challenge import (
SCOPE_RESPONSE_KEY,
OAuthScopeResponse,
finish_scope_response,
is_tool_call_request,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
ActingUser,
GrantedToolsetIds,
@ -604,7 +610,7 @@ if MCP_AVAILABLE:
name=LITELLM_MCP_SERVER_NAME,
version=LITELLM_MCP_SERVER_VERSION,
)
server.middleware.append(GatewayVersionPolicy())
server.middleware.extend((finish_scope_response, GatewayVersionPolicy()))
server.create_initialization_options = types.MethodType(_gateway_create_initialization_options, server)
sse: Final[SseServerTransport] = SseServerTransport("/sse/messages")
@ -2351,6 +2357,18 @@ if MCP_AVAILABLE:
await rejection(scope, receive, send)
return
scope_response: Final = (
OAuthScopeResponse(
send,
get_request_base_url(StarletteRequest(scope)),
get_route_relative_request_path(scope),
)
if is_tool_call_request(body) or len(body) >= _MCP_ROUTING_PEEK_MAX_BYTES
else None
)
if scope_response is not None:
scope[SCOPE_RESPONSE_KEY] = scope_response
use_stateful: Final = bool(session_id or is_initialize)
target_manager: Final = session_manager_stateful if use_stateful else session_manager_stateless
@ -2468,7 +2486,7 @@ if MCP_AVAILABLE:
touch_last_seen=(scope.get("method") or "").upper() != "DELETE",
copy_existing_session_auth_context=is_initialize,
)
local_send = send
local_send = scope_response or send
if use_stateful and is_initialize:
local_send = _wrap_send_with_stateful_session_auth_context(
local_send,

16
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -37831,6 +37831,22 @@
],
"title": "Resource"
}
},
{
"in": "query",
"name": "scope",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Scope"
}
}
],
"responses": {

3
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -1,4 +1,5 @@
import json
from collections.abc import Mapping
from datetime import datetime
from typing import Annotated, Any, Final, Literal
@ -29,6 +30,7 @@ MCPInfo = Annotated[dict[str, Any], AfterValidator(_validate_mcp_protocol_metada
class MCPOAuthMetadata(LiteLLMBaseModel):
authorization_server_metadata: Mapping[str, object] | None = Field(default=None, exclude=True)
scopes: list[str] | None = None
"""Resource-driven scopes for the authorization request: the RFC 9728 protected-resource
``scopes_supported``, or the ``scope`` from the WWW-Authenticate 401 challenge when the resource
@ -117,6 +119,7 @@ class MCPServer(LiteLLMBaseModel):
# scope=="user" values must be supplied per-user.
env_vars: list[dict[str, Any]] | None = None
# OAuth-specific fields
authorization_server_metadata: Mapping[str, object] | None = Field(default=None, exclude=True)
client_id: str | None = None
client_secret: str | None = None
issuer: str | None = None

148
tests/integration/mcp/test_mcp_oauth_flows.py
View file

@ -1,11 +1,13 @@
import base64
import hashlib
import json
import re
import secrets
import textwrap
import time
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final
@ -23,15 +25,19 @@ from integration._support.mcp import (
McpCaller,
McpPeer,
Outcome,
ScriptedTool,
_outcome_from_rpc,
call_tool,
mcp_peer,
register_mcp,
scripted_peer,
text_result,
tool_calls,
)
from integration._support.mcp_grants import create_toolset
from integration._support.oauth_server import AuthorizationServer, oauth_server
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
ADD: Final = {"a": 2, "b": 3}
CLIENT_REDIRECT: Final = "http://127.0.0.1:9/cb"
@ -240,7 +246,8 @@ class _Pkce:
def _authorize_through_gateway(
gateway: Gateway, auth: AuthorizationServer, alias: str, key: str, client_id: str, pkce: _Pkce
gateway: Gateway, auth: AuthorizationServer, alias: str, key: str, client_id: str, pkce: _Pkce,
scope: str = "tools.call",
) -> str:
started: Final = gateway.client.get(
f"/{alias}/authorize",
@ -251,7 +258,7 @@ def _authorize_through_gateway(
"state": "client-state",
"code_challenge": pkce.challenge,
"code_challenge_method": "S256",
"scope": "tools.call",
"scope": scope,
},
headers={"x-litellm-api-key": key},
)
@ -575,3 +582,140 @@ def test_token_exchange_callers_with_different_subject_tokens_own_separate_listi
assert second.list_tools().ok
assert _echoed_description(second.call(f"{alias}-add", probe)) == "Add two integers"
assert tool_calls(peer.drain()) == (), "a blocked probe reached the peer"
@contextmanager
def _scope_enforcing_peer(auth: AuthorizationServer) -> Iterator[tuple[McpPeer, Wire]]:
with scripted_peer(
ScriptedTool("read_op", lambda _: text_result("read completed")),
ScriptedTool("write_op", lambda _: text_result("write completed")),
ScriptedTool("forbid_op", lambda _: text_result("must not execute")),
) as inner:
def enforce(request: Request) -> Reply:
body: Final = json.loads(request.body) if request.body else {}
if body.get("method") == "tools/call":
name: Final = body["params"]["name"]
if name == "forbid_op":
return Reply(status=403, body=b"denied")
required: Final = "tools.write" if name == "write_op" else "tools.read"
token: Final = request.headers.get("authorization", "").removeprefix("Bearer ")
with auth.lock:
granted: Final = auth.access_tokens.get(token, {}).get("scope", "")
if required not in granted.split():
return Reply(status=403, headers={"WWW-Authenticate": f'Bearer error="insufficient_scope", scope="{required}"'})
forwarded: Final = httpx.request(request.method, inner.url, content=request.body, headers=dict(request.headers))
return Reply(status=forwarded.status_code, body=forwarded.content, content_type=forwarded.headers.get("content-type", "application/json"))
with wire_server(enforce) as wire:
yield McpPeer(wire.url + "/mcp", inner.calls), wire
@pytest.mark.parametrize("entry", ["rest", "server_mcp"])
def test_managed_scope_step_up_preserves_read_and_write_without_retry(gateway: Gateway, entry: EntryPoint) -> None:
with oauth_server(scopes=("tools.read", "tools.write")) as auth, _scope_enforcing_peer(auth) as (peer, wire), gateway.scenario() as scenario:
alias: Final = "scope" + uuid.uuid4().hex[:8]
identity: Final = _register_oauth(
scenario, peer, auth, alias, auth_type="oauth2", oauth2_flow="authorization_code",
credentials={"client_id": "scope-client", "client_secret": "scope-secret", "scopes": ["tools.read"]},
)
owner: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]})
stranger: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]})
pkce: Final = _Pkce(secrets.token_urlsafe(32))
code: Final = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.read")
first: Final = _redeem(gateway, alias, owner, "scope-client", code, pkce)
assert first.status_code == 200, first.text
caller: Final = McpCaller(gateway, owner, entry, alias)
read: Final = caller.call(f"{alias}-read_op", {}, identity)
assert read.ok and read.text == "read completed", read.raw
peer.drain()
wire.drain()
denied: Final = (
call_tool(gateway, owner, identity, f"{alias}-write_op", {}) if entry == "rest"
else caller.rpc("tools/call", {"name": f"{alias}-write_op", "arguments": {}})
)
assert denied.status_code == 403, denied.text
challenge: Final = denied.headers["www-authenticate"]
assert 'error="insufficient_scope"' in challenge and 'scope="tools.write"' in challenge
assert f'resource_metadata="{_base(gateway)}/.well-known/oauth-protected-resource/' in challenge
attempts: Final = tuple(request for request in wire.drain() if request.body and json.loads(request.body).get("method") == "tools/call")
assert len(attempts) == 1 and tool_calls(peer.drain()) == (), "denied operation was retried or executed"
second_code: Final = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.write")
second: Final = _redeem(gateway, alias, owner, "scope-client", second_code, pkce)
assert second.status_code == 200, second.text
with auth.lock:
granted: Final = auth.access_tokens[_issued_token(second.json())]["scope"]
assert granted.split() == ["tools.read", "tools.write"]
for operation in ("read", "write"):
result: Final = caller.call(f"{alias}-{operation}_op", {}, identity)
assert result.ok and result.text == f"{operation} completed", result.raw
forbidden: Final = caller.call(f"{alias}-forbid_op", {}, identity)
assert forbidden.status == 200 and not forbidden.ok, forbidden.raw
unauthorized: Final = McpCaller(gateway, stranger, entry, alias).call(f"{alias}-read_op", {}, identity)
assert unauthorized.status == 401, unauthorized.raw
@pytest.mark.parametrize("route", ["/mcp", "/toolset/{toolset}/mcp"])
def test_managed_step_up_advertises_sdk_compatible_metadata_and_bound_permissions(gateway: Gateway, route: str) -> None:
from mcp.shared.auth_utils import check_resource_allowed
with (
oauth_server(scopes=("tools.read", "tools.write")) as auth,
_scope_enforcing_peer(auth) as (peer, wire),
gateway.scenario() as scenario,
):
alias: Final = "scope" + uuid.uuid4().hex[:8]
identity: Final = _register_oauth(
scenario,
peer,
auth,
alias,
auth_type="oauth2",
oauth2_flow="authorization_code",
credentials={"client_id": "scope-client", "client_secret": "scope-secret", "scopes": ["tools.read"]},
)
toolset_name: Final = "scope" + uuid.uuid4().hex[:8]
toolset_id: Final = create_toolset(
scenario, ((identity, "read_op"), (identity, "write_op")), toolset_name=toolset_name
)
owner: Final = scenario.key(
user_id=scenario.user(), object_permission={"mcp_servers": [identity], "mcp_toolsets": [toolset_id]}
)
pkce: Final = _Pkce(secrets.token_urlsafe(32))
code: Final = _authorize_through_gateway(gateway, auth, alias, owner, "scope-client", pkce, "tools.read")
first: Final = _redeem(gateway, alias, owner, "scope-client", code, pkce)
assert first.status_code == 200, first.text
path: Final = route.format(toolset=toolset_name)
headers: Final = {**ACCEPT, "Authorization": f"Bearer {owner}"}
initialized: Final = gateway.client.post(
path, headers=headers, json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": dict(INITIALIZE)}
)
assert initialized.status_code == 200, initialized.text
peer.drain()
wire.drain()
response: Final = gateway.client.post(
path,
headers=headers,
json={
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": {"name": f"{alias}-write_op", "arguments": {}},
},
)
assert response.status_code == 403, response.text
challenge: Final = response.headers["www-authenticate"]
metadata_url: Final = challenge.split('resource_metadata="', 1)[1].split('"', 1)[0]
metadata: Final = gateway.client.get(metadata_url)
assert metadata.status_code == 200, metadata.text
assert check_resource_allowed(
requested_resource=_base(gateway) + path, configured_resource=metadata.json()["resource"]
)
encoded_scope: Final = challenge.split('scope="', 1)[1].split('"', 1)[0]
assert encoded_scope == f"litellm:mcp_scope:{identity.encode().hex()}:{b'tools.write'.hex()}"
assert tool_calls(peer.drain()) == (), "denied operation executed upstream"
attempts: Final = tuple(
request
for request in wire.drain()
if request.body and json.loads(request.body).get("method") == "tools/call"
)
assert len(attempts) == 1, "denied operation was retried"

76
tests/unit/experimental_mcp_client/test_mcp_client.py
View file

@ -2757,6 +2757,35 @@ async def test_optional_discovery_allows_exhaustion_at_page_cap(method: str, mon
)
def test_standalone_client_request_does_not_require_proxy_dependencies():
import textwrap
from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter
script = textwrap.dedent("""
import asyncio
import sys
for name in ("fastapi", "orjson", "email_validator", "annotated_doc"):
sys.modules[name] = None
import httpx2
from litellm.experimental_mcp_client.client import MCPClient
async def request():
upstream = MCPClient(server_url="https://upstream.example/mcp")
transport = httpx2.MockTransport(lambda request: httpx2.Response(200))
async with upstream._create_httpx_client_factory(transport=transport)() as client:
response = await client.post(upstream.server_url, json={"method": "tools/call"})
assert response.status_code == 200
assert all(sys.modules[name] is None for name in ("fastapi", "orjson", "email_validator", "annotated_doc"))
asyncio.run(request())
print("standalone request succeeded")
""")
result = run_child_interpreter(script, timeout=60)
assert result.returncode == 0, result.stderr
assert result.stdout.strip() == "standalone request succeeded"
def test_client_import_before_proxy_credentials_succeeds_in_fresh_process():
import subprocess
@ -4942,3 +4971,50 @@ async def test_persistent_session_ends_after_a_broken_stream_so_the_next_call_ge
await asyncio.wait_for(session.run(broken_stream), 5)
await asyncio.wait_for(session.wait_closed(), 5)
assert session.closed, "a dead transport must end the session instead of being reused"
@pytest.mark.asyncio
@pytest.mark.parametrize("session_started_in_request", [True, False])
async def test_persistent_session_releases_the_current_requests_scope_response(session_started_in_request) -> None:
from mcp.server.context import ServerRequestContext
from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse
app = _stateful_upstream()
async with app.router.lifespan_context(app):
client = _StatefulUpstreamClient(app, server_url="http://upstream/mcp", transport_type=MCPTransport.http)
first = OAuthScopeResponse(AsyncMock(), "https://gateway.example", "/mcp")
second = OAuthScopeResponse(AsyncMock(), "https://gateway.example", "/mcp")
context = ServerRequestContext(
session=MagicMock(),
lifespan_context={},
protocol_version="2025-11-25",
method="tools/call",
request=Request({"type": "http", SCOPE_RESPONSE_KEY: first}),
)
earlier_session = None if session_started_in_request else client.open_persistent_session()
token = active_mcp_request_ctx_var.set(context)
session = earlier_session if earlier_session is not None else client.open_persistent_session()
try:
selected = await client.call_tool(
CallToolRequestParams(name="select_project", arguments={"name": "demo"}), persistent_session=session
)
assert not selected.is_error and not first.pending
active_mcp_request_ctx_var.set(
ServerRequestContext(
session=context.session,
lifespan_context={},
protocol_version="2025-11-25",
method="tools/call",
request=Request({"type": "http", SCOPE_RESPONSE_KEY: second}),
)
)
created = await client.call_tool(
CallToolRequestParams(name="create_feature", arguments={"title": "feature"}), persistent_session=session
)
assert not created.is_error and created.content[0].text == "demo/feature"
assert not second.pending
finally:
active_mcp_request_ctx_var.reset(token)
session.close()
await asyncio.wait_for(session.wait_closed(), 5)

29
tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py
View file

@ -267,3 +267,32 @@ def test_pure_non_auth_response_still_classifies_upstream_error():
fault = classify_list_exception(exc)
assert fault.tag == "upstream_error"
assert fault.status_code == 502
@pytest.mark.parametrize("status, challenge, expected", [
(403, 'Bearer error="insufficient_scope", scope="tools.write"', "tools.write"),
(403, 'bEaReR error=insufficient_scope, scope="tools.read tools.write"', "tools.read tools.write"),
(403, 'Bearer error="insufficient_scope"', ""),
(403, 'Bearer error="insufficient_scope", scope=', None),
(403, 'Bearer error="insufficient_scope", scope', None),
(401, 'Bearer error="insufficient_scope", scope="tools.write"', None),
(403, 'Bearer error="invalid_token", scope="tools.write"', None),
(403, '', None),
(403, 'Basic error="insufficient_scope", scope="tools.write"', None),
(403, 'Bearer error_description="error=insufficient_scope", scope="tools.write"', None),
(403, 'Bearer error="insufficient_scope", error="invalid_token"', None),
(403, 'Bearer error="insufficient_scope", scope="tools.write\r\ninjected"', None),
(403, 'Bearer error="insufficient_scope", scope="tools.write", Basic realm="other"', None),
])
@pytest.mark.parametrize("client", ["httpx", "httpx2"])
def test_managed_scope_step_up_only_accepts_unambiguous_bearer_challenges(status, challenge, expected, client):
import httpx2
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import upstream_insufficient_scope
transport = httpx if client == "httpx" else httpx2
response = transport.Response(status, headers={"www-authenticate": challenge}, request=transport.Request("POST", "https://peer.test/mcp"))
cause = transport.HTTPStatusError("upstream rejected", request=response.request, response=response)
wrapped = RuntimeError("transport failure")
wrapped.__cause__ = cause
assert upstream_insufficient_scope(wrapped) == expected

20
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py
View file

@ -210,16 +210,17 @@ async def test_unrotated_refresh_token_is_carried_forward():
@pytest.mark.asyncio
async def test_unrecorded_scope_is_carried_forward():
@pytest.mark.parametrize("prior_scopes", [(), ("read", "write")])
async def test_unrecorded_scope_is_carried_forward(prior_scopes):
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"))
"a", "s", OAuthToken("old", refresh_token="rt", scopes=prior_scopes)
)
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")
assert token.scopes == prior_scopes
assert persisted[0][5] == prior_scopes
@pytest.mark.asyncio
@ -257,18 +258,19 @@ async def test_refresh_omits_resource_when_unset():
@pytest.mark.asyncio
async def test_returned_scope_overrides_prior_when_present():
@pytest.mark.parametrize("scope, expected", [("read", ("read",)), ("", ()), (" ", ())])
async def test_returned_scope_overrides_prior_when_present(scope, expected):
persisted = []
refresher = _refresher(
body={"access_token": "new-at", "scope": "read"},
body={"access_token": "new-at", "scope": scope},
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",)
assert token.scopes == expected
assert persisted[0][5] == expected
@pytest.mark.asyncio
@ -366,7 +368,7 @@ async def test_cimd_refresh_on_fresh_replica_preserves_user_and_client_identity(
{},
)
]
assert persisted == [("alice", "srv", "new-at", "old-rt", 3600, None)]
assert persisted == [("alice", "srv", "new-at", "old-rt", 3600, ())]
assert server.client_id is None

57
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py
View file

@ -26,6 +26,63 @@ class _RecordingStore:
self.invalidations.append((user_id, server_id))
@pytest.fixture
def credential_database(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy import proxy_server
prisma = MagicMock()
table = prisma.db.litellm_mcpusercredentials
table.upsert = AsyncMock()
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setenv("LITELLM_SALT_KEY", "credential-persistence-test-salt")
return table
@pytest.mark.asyncio
@pytest.mark.parametrize("scopes, expected", [(None, None), ((), []), (("read",), ["read"])])
async def test_credential_persistence_preserves_known_and_unknown_scopes(credential_database, scopes, expected):
from litellm.proxy._experimental.mcp_server.db import decode_oauth_payload
from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import _persist_credential
await _persist_credential("alice", "srv", "access", "refresh", 60, scopes)
row = credential_database.upsert.call_args.kwargs["data"]["create"]
payload = decode_oauth_payload(row["credential_b64"])
assert payload is not None and payload.get("scopes") == expected
assert row["user_id"] == "alice" and row["server_id"] == "srv"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prior_scopes, response_scope, expected",
[((), None, []), (("read", "write"), "", []), (("read", "write"), None, ["read", "write"])],
)
async def test_native_refresh_persists_actual_grant_through_storage_adapter(
credential_database, prior_scopes, response_scope, expected,
):
from litellm.proxy._experimental.mcp_server.db import decode_oauth_payload
from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import AuthorizationCodeRefresher
from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import _persist_credential
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="srv", name="srv", transport="http", auth_type="oauth2",
token_url="https://idp.example/token", client_id="client", client_secret="secret",
)
async def token_endpoint(url, form, headers):
assert url == server.token_url and form["grant_type"] == "refresh_token"
return {"access_token": "rotated", **({"scope": response_scope} if response_scope is not None else {})}
refresher = AuthorizationCodeRefresher(lambda _: server, token_endpoint, _persist_credential)
refreshed = await refresher.refresh("alice", "srv", OAuthToken("old", refresh_token="refresh", scopes=prior_scopes))
assert refreshed is not None and refreshed.scopes == tuple(expected)
payload = decode_oauth_payload(credential_database.upsert.call_args.kwargs["data"]["create"]["credential_b64"])
assert payload is not None and payload.get("scopes") == expected
assert payload["access_token"] == "rotated"
class _BlockingStore:
def __init__(self, access_token: str) -> None:
self._access_token = access_token

12
tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
View file

@ -1856,3 +1856,15 @@ async def test_saved_cimd_grant_refreshes_from_encrypted_storage_on_a_fresh_repl
assert parse_qs(upstream.calls[0].request.content.decode())["client_id"] == [identity]
assert "client_secret" not in parse_qs(upstream.calls[0].request.content.decode())
assert server.client_id is None and server.client_id_metadata_document_supported is capability
@pytest.mark.asyncio
@pytest.mark.parametrize("provider_scope, expected", [(None, ["a"]), ("b", ["b"]), ("", [])])
async def test_silent_refresh_preserves_omitted_scope_but_accepts_explicit_grant(monkeypatch, provider_scope, expected):
import litellm.proxy._experimental.mcp_server.db as db_mod
await _run_refresh(monkeypatch, _refresh_server(), {
"access_token": "new-access",
**({"scope": provider_scope} if provider_scope is not None else {}),
})
assert db_mod.store_user_oauth_credential.await_args.kwargs["scopes"] == expected

207
tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -11981,6 +11981,7 @@ def jwt_oauth_identity(monkeypatch: pytest.MonkeyPatch) -> tuple["JWTHandler", "
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
prisma: Final = MagicMock()
prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=None)
prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
return handler, signing_key
@ -13732,7 +13733,7 @@ async def test_token_route_preserves_saved_cimd_grant_after_origin_change(
"matching", "unicode_grant", "foreign_refresh", "missing_refresh", "missing_grant", "database_missing",
"database_outage", "static_client", "anonymous",
])
async def test_saved_cimd_refresh_identity_is_bound_to_the_callers_stored_grant(
async def test_saved_refresh_metadata_is_bound_to_the_callers_stored_grant(
monkeypatch: pytest.MonkeyPatch, state: str,
) -> None:
from types import SimpleNamespace
@ -13766,12 +13767,210 @@ async def test_saved_cimd_refresh_identity_is_bound_to_the_callers_stored_grant(
table.find_unique.side_effect = RuntimeError("database unavailable")
if state == "static_client":
server.client_id = "configured-client"
resolved: Final = await endpoints._saved_cimd_refresh_client_id(
resolved: Final = await endpoints._saved_refresh_credential(
server, None if state == "anonymous" else "alice",
"foreign-refresh" if state == "foreign_refresh" else grant,
)
assert resolved == (identity if state in ("matching", "unicode_grant") else None)
if state in ("anonymous", "static_client", "database_missing"):
if state in ("matching", "unicode_grant", "static_client"):
assert resolved is not None and resolved["cimd_client_id"] == identity
assert resolved["refresh_token"] == grant
else:
assert resolved is None
if state in ("anonymous", "database_missing"):
table.find_unique.assert_not_awaited()
else:
table.find_unique.assert_awaited_once_with(where={"user_id_server_id": {"user_id": "alice", "server_id": server.server_id}})
@pytest.mark.asyncio
@pytest.mark.parametrize("prior_scopes, expected", [(["tools.read"], "tools.read tools.write"), (None, "tools.write"), ([], "tools.write")])
async def test_managed_authorize_retains_the_callers_previous_grant(jwt_oauth_identity, monkeypatch, prior_scopes, expected):
from types import SimpleNamespace
from urllib.parse import parse_qs, urlsplit
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
_, signing_key = jwt_oauth_identity
monkeypatch.setenv("LITELLM_SALT_KEY", "scope-grant-test-salt")
server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client", "scopes": ["tools.read"]})
global_mcp_server_manager.registry[server.server_id] = server
credential = db.encrypt_value_helper(json.dumps({"type": "oauth2", "access_token": "old-access", "scopes": prior_scopes}))
table = proxy_server.prisma_client.db.litellm_mcpusercredentials
table.find_unique = AsyncMock(return_value=SimpleNamespace(credential_b64=credential))
request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"})
response = await endpoints.authorize_with_server(
request, server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.write",
code_challenge="challenge", code_challenge_method="S256",
)
assert parse_qs(urlsplit(response.headers["location"]).query)["scope"] == [expected]
assert table.find_unique.call_args.kwargs["where"] == {
"user_id_server_id": {"user_id": "jwt-owner", "server_id": server.server_id}
}
@pytest.mark.asyncio
@pytest.mark.parametrize("returned_scope, expected", [(None, ["tools.read", "tools.write"]), ("tools.read", ["tools.read"]), ("", [])])
async def test_managed_code_exchange_preserves_requested_scope_only_when_omitted(
jwt_oauth_identity, monkeypatch, respx_mock, returned_scope, expected,
):
from http.cookies import SimpleCookie
from urllib.parse import parse_qs, urlsplit
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
_, signing_key = jwt_oauth_identity
monkeypatch.setenv("LITELLM_SALT_KEY", "requested-scope-test-salt")
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"})
global_mcp_server_manager.registry[server.server_id] = server
table = proxy_server.prisma_client.db.litellm_mcpusercredentials
table.find_unique = AsyncMock(return_value=None)
table.upsert = AsyncMock()
request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"})
started = await endpoints.authorize_with_server(
request, server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.read tools.write",
code_challenge="challenge", code_challenge_method="S256",
)
state = parse_qs(urlsplit(started.headers["location"]).query)["state"][0]
cookies = SimpleCookie()
cookies.load(started.headers["set-cookie"])
callback_request = _token_request({"Cookie": "; ".join(f"{name}={value.value}" for name, value in cookies.items())})
returned = await endpoints.callback(callback_request, code="provider-code", state=state, error=None, iss=None)
code = parse_qs(urlsplit(returned.headers["location"]).query)["code"][0]
upstream = respx_mock.post(server.token_url).respond(200, json={
"access_token": "new-access", "token_type": "Bearer", "refresh_token": "new-refresh",
**({"scope": returned_scope} if returned_scope is not None else {}),
})
exchanged = await endpoints.exchange_token_with_server(
request, server, "authorization_code", client_id="scope-client", code=code, code_verifier="verifier", redirect_uri="http://127.0.0.1:9/cb", client_secret=None,
)
assert parse_qs(upstream.calls[0].request.content.decode())["code"] == ["provider-code"]
saved = db.decode_oauth_payload(table.upsert.call_args.kwargs["data"]["create"]["credential_b64"])
assert saved is not None and saved.get("scopes") == expected
assert json.loads(exchanged.body)["scope"] == " ".join(expected or ())
@pytest.mark.asyncio
@pytest.mark.parametrize("advertised, expected", [(True, "tools.read offline_access"), (False, "tools.read")])
async def test_managed_authorize_requests_only_selected_scope_and_supported_refresh_access(monkeypatch, advertised, expected):
from urllib.parse import parse_qs, urlsplit
from mcp.shared.auth import OAuthMetadata
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
monkeypatch.setenv("LITELLM_SALT_KEY", "metadata-scope-test-salt")
metadata = OAuthMetadata(
issuer="https://idp.example.com", authorization_endpoint="https://idp.example.com/authorize",
token_endpoint="https://idp.example.com/token", response_types_supported=["code"],
scopes_supported=["tools.read", "admin"] + (["offline_access"] if advertised else []),
)
server = _cimd_oauth_server().model_copy(update={"authorization_server_metadata": metadata.model_dump(mode="json")})
response = await endpoints.authorize_with_server(
_cimd_request(), server, "scope-client", "http://127.0.0.1:9/cb", scope="tools.read",
code_challenge="challenge", code_challenge_method="S256",
)
assert parse_qs(urlsplit(response.headers["location"]).query)["scope"] == [expected]
@pytest.mark.asyncio
@pytest.mark.parametrize("mismatch", ["server", "issuer", "user", "mode", "tampered", "invalid_payload"])
async def test_managed_scope_code_cannot_cross_oauth_boundaries(jwt_oauth_identity, monkeypatch, respx_mock, mismatch):
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
_, signing_key = jwt_oauth_identity
monkeypatch.setenv("LITELLM_SALT_KEY", "scope-boundary-test-salt")
server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"})
payload = endpoints._ManagedAuthorizationCode(
upstream_code="provider-code", mcp_server_id="other-server" if mismatch == "server" else server.server_id,
issuer="https://other-issuer.test" if mismatch == "issuer" else server.issuer,
requested_scope="tools.read tools.write", scope_user_id="other-user" if mismatch == "user" else "jwt-owner",
)
code = endpoints._MANAGED_AUTH_CODE_PREFIX + ("invalid" if mismatch == "tampered" else endpoints.encrypt_value_helper("{}" if mismatch == "invalid_payload" else payload.model_dump_json()))
request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"})
if mismatch == "mode":
server.auth_type = MCPAuth.true_passthrough
with pytest.raises(HTTPException) as raised:
await endpoints.exchange_token_with_server(
request, server, "authorization_code", code=code, redirect_uri="http://127.0.0.1:9/cb",
client_id="scope-client", client_secret=None, code_verifier="verifier",
)
assert raised.value.status_code == 400
assert respx_mock.calls.call_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("scope, requested_scope, expected", [
(None, None, ["tools.read", "tools.write"]), ("tools.read", None, ["tools.read"]), ("", None, []),
(None, "tools.write", ["tools.write"]), ("tools.read", "tools.write", ["tools.read"]), ("", "tools.write", []),
])
async def test_managed_refresh_retains_only_the_actual_grant(jwt_oauth_identity, monkeypatch, respx_mock, scope, requested_scope, expected):
from types import SimpleNamespace
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
_, signing_key = jwt_oauth_identity
monkeypatch.setenv("LITELLM_SALT_KEY", "refresh-scope-test-salt")
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
server = _cimd_oauth_server().model_copy(update={"client_id": "scope-client"})
global_mcp_server_manager.registry[server.server_id] = server
credential = db.encrypt_value_helper(json.dumps({
"type": "oauth2", "access_token": "old-access", "refresh_token": "owner-refresh", "scopes": ["tools.read", "tools.write"],
}))
table = proxy_server.prisma_client.db.litellm_mcpusercredentials
table.find_unique = AsyncMock(return_value=SimpleNamespace(credential_b64=credential))
table.upsert = AsyncMock()
upstream = respx_mock.post(server.token_url).respond(200, json={
"access_token": "new-access", "token_type": "Bearer", "refresh_token": "new-refresh",
**({"scope": scope} if scope is not None else {}),
})
request = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, scope='litellm_proxy_admin')}"})
await endpoints.exchange_token_with_server(
request, server, "refresh_token", code=None, redirect_uri=None, client_id="scope-client", client_secret=None,
code_verifier=None, refresh_token="owner-refresh", scope=requested_scope,
)
saved = db.decode_oauth_payload(table.upsert.call_args.kwargs["data"]["create"]["credential_b64"])
assert saved is not None and saved.get("scopes") == expected
assert upstream.call_count == 1
from urllib.parse import parse_qs
sent = parse_qs(upstream.calls.last.request.content.decode())
assert sent.get("scope") == ([requested_scope] if requested_scope else None)
@pytest.mark.asyncio
@pytest.mark.parametrize("per_server_relay", [False, True])
async def test_managed_resource_metadata_excludes_refresh_permission(monkeypatch, per_server_relay):
from litellm.proxy._experimental.mcp_server import discoverable_endpoints as endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
server = _cimd_oauth_server().model_copy(update={
"scopes": ["tools.read", "offline_access"], "per_server_oauth_discovery": per_server_relay,
})
global_mcp_server_manager.registry[server.server_id] = server
result = await endpoints._build_oauth_protected_resource_response(_cimd_request(), server.server_id, True)
assert result["scopes_supported"] == ["tools.read"]
@pytest.mark.asyncio
@pytest.mark.parametrize("granted, required, expected", [
(["tools.read"], ("tools.write",), "absent"),
(["tools.read", "tools.write"], ("tools.write",), "present"),
(None, ("tools.write",), "absent"),
(["tools.read"], (), "present"),
])
async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, required, expected):
from litellm.proxy._experimental.mcp_server import db, discoverable_endpoints as endpoints
credential = {**_stored_grant(), "scopes": granted}
read = AsyncMock(return_value=credential)
prisma = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
monkeypatch.setattr(db, "get_user_oauth_credential", read)
assert await endpoints._vendor_credential_state("owner", "server", required) == expected
read.assert_awaited_once_with(prisma, "owner", "server")

297
tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
View file

@ -991,7 +991,7 @@ SCOPED_RESOURCE = "https://llm.example.com/mcp/github"
_MANAGER_PATCH = "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager"
def _scoped_authorize(client_id, resource, session_user_id="u1"):
def _scoped_authorize(client_id, resource, session_user_id="u1", **kwargs):
return aggregate_authorize(
request=_request(query=f"client_id={client_id}"),
client_id=client_id,
@ -1002,6 +1002,7 @@ def _scoped_authorize(client_id, resource, session_user_id="u1"):
response_type="code",
session_user_id=session_user_id,
resource=resource,
**kwargs,
)
@ -1914,7 +1915,9 @@ def _redis_that(async_increment, get=None):
cache.redis_cache.async_increment = async_increment
cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key)
cache.redis_cache.init_async_client.return_value.get = get or AsyncMock(return_value=None)
cache.redis_cache.async_get_cache = AsyncMock(side_effect=AssertionError("peek must read the client, not the wrapper"))
cache.redis_cache.async_get_cache = AsyncMock(
side_effect=AssertionError("peek must read the client, not the wrapper")
)
cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory"))
return cache
@ -2068,7 +2071,11 @@ async def test_refresh_of_a_rotated_token_answers_503_before_minting_while_redis
(await _redeem_native(await _native_code(client_id, cache=issued), client_id, _Minter(), cache=issued)).body
)
rotated = json.loads(
(await _refresh_native(payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)))).body
(
await _refresh_native(
payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1))
)
).body
)["refresh_token"]
minter = _Minter()
@ -2578,3 +2585,287 @@ async def test_token_exchange_relays_a_mint_refusal(failure, status, error):
response = await _exchange_native(client_id, _Minter(failure), _Exchanger())
assert response.status_code == status
assert json.loads(response.body)["error"] == error
@pytest.mark.asyncio
async def test_connect_flow_retains_step_up_scope_until_the_vendor_grants_it():
from unittest.mock import AsyncMock, patch
client_id = (await _register([REDIRECT_URI]))["client_id"]
server = _scoped_mcp_server()
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = server
response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope="tools.write")
vendor = AsyncMock(return_value="absent")
described = await _describe_page(response, scoped_server=server, vendor=vendor)
assert json.loads(described.body)["requested_scopes"] == ["tools.write"]
assert json.loads(described.body)["connected"] is False
vendor.assert_awaited_once_with("u1", server.server_id, ("tools.write",))
premature = await _complete_page(response, scoped_server=server, vendor=vendor)
assert premature.status_code == 400
vendor.return_value = "present"
complete = await _complete_page(response, scoped_server=server, vendor=vendor)
assert complete.status_code == 303
assert parse_qs(urlparse(complete.headers["location"]).query)["code"][0].startswith(GATEWAY_AUTH_CODE_PREFIX)
@pytest.mark.asyncio
async def test_aggregate_step_up_requires_the_upstream_grant_without_narrowing_gateway_access():
from unittest.mock import AsyncMock
client_id = (await _register([REDIRECT_URI]))["client_id"]
server = _scoped_mcp_server()
encoded = f"litellm:mcp_scope:{server.server_id.encode().hex()}:{b'tools.write'.hex()}"
response = _scoped_authorize(client_id, "https://llm.example.com/mcp", scope=encoded)
vendor = AsyncMock(return_value="absent")
described = await _describe_page(response, scoped_server=server, vendor=vendor)
assert json.loads(described.body)["connected"] is False
assert json.loads(described.body)["requested_scopes"] == ["tools.write"]
premature = await _complete_page(response, scoped_server=server, vendor=vendor)
assert premature.status_code == 400
vendor.return_value = "present"
ready = await _describe_page(response, scoped_server=server, vendor=vendor)
assert json.loads(ready.body)["state"] == "unscoped"
assert json.loads(ready.body)["server_id"] is None
completed = await _complete_page(response, scoped_server=server, vendor=vendor)
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
tokens = await _redeem(code, client_id)
assert tokens.status_code == 200
assert _opened_principal(json.loads(tokens.body)).resource_server_id is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
"scope",
[
"litellm:mcp_scope:",
"litellm:mcp_scope:aa:bb",
"litellm:mcp_scope:6769746875622d6964:20",
"litellm:mcp_scope:6964:5A",
],
)
async def test_aggregate_step_up_rejects_malformed_permissions(scope):
client_id = (await _register([REDIRECT_URI]))["client_id"]
response = _scoped_authorize(client_id, "https://llm.example.com/mcp", scope=scope)
assert response.status_code == 400
assert json.loads(response.body)["error"] == "invalid_scope"
@pytest.mark.asyncio
async def test_step_up_validates_every_target_before_credential_reads_and_preserves_each_grant():
from unittest.mock import AsyncMock, patch
from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope
client_id = (await _register([REDIRECT_URI]))["client_id"]
first, second = _scoped_mcp_server("first"), _scoped_mcp_server("second")
scopes = " ".join(
[encode_upstream_scope(first.server_id, "first.write"), encode_upstream_scope(second.server_id, "second.write")]
)
response = _scoped_authorize(client_id, "https://llm.example.com", scope=scopes)
handle, cookies = _flow_cookie_from(response)
vendor = AsyncMock(return_value="absent")
reachable = AsyncMock(side_effect=[True, False])
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_id.side_effect = lambda identity: {
first.server_id: first,
second.server_id: second,
}.get(identity)
refused = await describe_connect_flow(_request(cookies=cookies), handle, "u1", vendor, reachable)
assert refused.status_code == 400
vendor.assert_not_awaited()
reachable.side_effect = None
reachable.return_value = True
vendor.side_effect = ["present", "absent"]
pending = await describe_connect_flow(_request(cookies=cookies), handle, "u1", vendor, reachable)
assert json.loads(pending.body)["server_id"] == second.server_id
assert json.loads(pending.body)["requested_scopes"] == ["second.write"]
assert [call.args for call in vendor.await_args_list] == [
("u1", first.server_id, ("first.write",)),
("u1", second.server_id, ("second.write",)),
]
vendor.reset_mock()
wrong_user = await describe_connect_flow(_request(cookies=cookies), handle, "u2", vendor, reachable)
assert wrong_user.status_code == 403
vendor.assert_not_awaited()
vendor.side_effect = None
vendor.return_value = "present"
complete = await complete_connect_flow(
_request(cookies=cookies, method="POST"),
handle,
"u1",
DualCache(),
lookup_vendor_credential=vendor,
lookup_server_reachability=reachable,
)
assert complete.status_code == 303
assert vendor.await_count == 2
@pytest.mark.asyncio
async def test_named_resource_rejects_another_upstreams_permission_before_consent():
from unittest.mock import patch
from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope
client_id = (await _register([REDIRECT_URI]))["client_id"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = _scoped_mcp_server()
response = _scoped_authorize(client_id, SCOPED_RESOURCE, scope=encode_upstream_scope("other-id", "write"))
assert response.status_code == 400
assert json.loads(response.body)["error"] == "invalid_scope"
@pytest.mark.parametrize(
"identity, permission",
[
("server:id/one", "tools.write"),
("snowman-☃", "https://permissions.example/read"),
("id", "!#$%&'()*+,-./:;<=>?@[]^_`{|}~"),
],
)
def test_upstream_permission_scope_round_trip_preserves_identity_and_permission(identity, permission):
from litellm.proxy._experimental.mcp_server.oauth_utils import decode_upstream_scope, encode_upstream_scope
encoded = encode_upstream_scope(identity, permission)
decoded = decode_upstream_scope(encoded)
assert decoded is not None
assert (decoded.server_id, decoded.permission) == (identity, permission)
assert " " not in encoded
@pytest.mark.asyncio
@pytest.mark.parametrize("state, expected", [("absent", 400), ("unavailable", 503)])
async def test_upstream_step_up_cannot_finish_with_an_unproven_grant(state, expected):
from unittest.mock import AsyncMock
from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope
client_id = (await _register([REDIRECT_URI]))["client_id"]
server = _scoped_mcp_server()
response = _scoped_authorize(
client_id, "https://llm.example.com/mcp", scope=encode_upstream_scope(server.server_id, "tools.write")
)
vendor = AsyncMock(return_value=state)
completed = await _complete_page(response, scoped_server=server, vendor=vendor)
assert completed.status_code == expected
assert "location" not in completed.headers
vendor.return_value = "present"
retried = await _complete_page(response, scoped_server=server, vendor=vendor)
assert retried.status_code == 303
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ["/mcp", "/toolset/analysis/mcp"])
async def test_sdk_completes_aggregate_step_up_discovery_consent_and_pkce(route):
from unittest.mock import AsyncMock
import httpx2
from mcp.client.auth import OAuthClientProvider
from mcp.shared.auth import AuthorizationCodeResult, OAuthClientMetadata
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_build_aggregate_authorization_server_response,
oauth_protected_resource_root,
)
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
server = _scoped_mcp_server()
storage = AsyncMock()
storage.get_tokens.return_value = None
storage.get_client_info.return_value = None
callback = AsyncMock()
vendor = AsyncMock(return_value="absent")
cache = DualCache()
async def consent(url):
params = {key: values[0] for key, values in parse_qs(urlparse(url).query).items()}
response = aggregate_authorize(request=_request(), session_user_id="u1", **params)
described = await _describe_page(response, scoped_server=server, vendor=vendor)
assert json.loads(described.body)["requested_scopes"] == ["tools.write"]
assert json.loads(described.body)["connected"] is False
premature = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache)
assert premature.status_code == 400
vendor.return_value = "present"
completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache)
assert completed.status_code == 303
returned = parse_qs(urlparse(completed.headers["location"]).query)
callback.return_value = AuthorizationCodeResult(code=returned["code"][0], state=returned["state"][0])
provider = OAuthClientProvider(
server_url=f"https://llm.example.com{route}",
client_metadata=OAuthClientMetadata(redirect_uris=[REDIRECT_URI], token_endpoint_auth_method="none"),
storage=storage,
redirect_handler=consent,
callback_handler=callback,
)
flow = provider.async_auth_flow(httpx2.Request("POST", f"https://llm.example.com{route}"))
initial = await anext(flow)
challenge = MCPUpstreamAuthError(
status_code=403,
www_authenticate=None,
server_name=server.server_name,
server_id=server.server_id,
required_scope="tools.write",
).to_http_exception(base_url="https://llm.example.com", request_path=route)
request = await flow.asend(httpx2.Response(challenge.status_code, headers=challenge.headers, request=initial))
assert request.url.path == "/.well-known/oauth-protected-resource"
request = await flow.asend(httpx2.Response(200, json=oauth_protected_resource_root(_request()), request=request))
assert request.url.path == "/.well-known/oauth-authorization-server/mcp"
request = await flow.asend(
httpx2.Response(200, json=_build_aggregate_authorization_server_response(_request(), False), request=request)
)
assert request.url.path == "/register"
registered = await register_aggregate_client(
_request("/register", method="POST"), json.loads(request.content), token_exchange_available=False
)
request = await flow.asend(httpx2.Response(registered.status_code, content=registered.body, request=request))
assert request.url.path == "/token"
params = {key: values[0] for key, values in parse_qs(request.content.decode()).items()}
tokens = await _redeem(cache=cache, **params)
assert tokens.status_code == 200
payload = json.loads(tokens.body)
assert _opened_principal(payload).resource_server_id is None
retry = await flow.asend(httpx2.Response(200, json=payload, request=request))
assert retry.url == initial.url
assert retry.headers["authorization"] == f"Bearer {payload['access_token']}"
with pytest.raises(StopAsyncIteration):
await flow.asend(httpx2.Response(200, json={"result": "write completed"}, request=retry))
callback.assert_awaited_once()
storage.set_tokens.assert_awaited_once()
@pytest.mark.parametrize(
"scope",
["tools.write", "litellm:mcp_scope::7772697465", "litellm:mcp_scope:6964:5A", "litellm:mcp_scope:6964:7772 697465"],
)
def test_upstream_permission_rejects_noncanonical_tokens(scope):
from litellm.proxy._experimental.mcp_server.oauth_utils import decode_upstream_scope
assert decode_upstream_scope(scope) is None
@pytest.mark.asyncio
async def test_named_step_up_finish_describes_and_retains_its_single_server_ceiling():
from unittest.mock import AsyncMock, patch
from litellm.proxy._experimental.mcp_server.oauth_utils import encode_upstream_scope
client_id = (await _register([REDIRECT_URI]))["client_id"]
server = _scoped_mcp_server()
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = server
response = _scoped_authorize(
client_id, SCOPED_RESOURCE, scope=encode_upstream_scope(server.server_id, "tools.write")
)
vendor = AsyncMock(return_value="present")
ready = await _describe_page(response, scoped_server=server, vendor=vendor)
payload = json.loads(ready.body)
assert payload["state"] == "interactive"
assert payload["server_id"] == server.server_id
assert payload["connected"] is True
completed = await _complete_page(response, scoped_server=server, vendor=vendor)
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
tokens = await _redeem(code, client_id)
assert tokens.status_code == 200
assert _opened_principal(json.loads(tokens.body)).resource_server_id == server.server_id

18
tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py
View file

@ -166,7 +166,11 @@ async def test_persistent_tool_call_with_limit_one_completes_without_nested_admi
oauth2_headers: Final = (
{"Authorization": "Bearer subject-token"} if auth_type == MCPAuth.oauth2_token_exchange else None
)
context_token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object()))
context_token: Final = active_mcp_request_ctx_var.set(
ServerRequestContext(
session=AsyncMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None
)
)
manager.track_gateway_session(gateway_session_id)
try:
with patch.object(manager, "create_mcp_client", new=AsyncMock(return_value=client)):
@ -209,7 +213,11 @@ async def test_persistent_session_admission_uses_updated_server_limit() -> None:
submitted=second_submitted,
release=release_second,
)
context_token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object()))
context_token: Final = active_mcp_request_ctx_var.set(
ServerRequestContext(
session=AsyncMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None
)
)
manager.track_gateway_session(gateway_session_id)
async def dispatch() -> CallToolResult:
@ -275,7 +283,11 @@ async def test_queued_persistent_call_survives_header_change_while_waiting_for_a
second_client: Final = _AdmissionProbeClient(
server.server_id, "second-fingerprint", tracker, submitted=second_submitted
)
context_token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object()))
context_token: Final = active_mcp_request_ctx_var.set(
ServerRequestContext(
session=AsyncMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None
)
)
manager.track_gateway_session(gateway_session_id)
async def hold_slot() -> None:

243
tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py
View file

@ -1,6 +1,8 @@
from litellm.proxy._experimental.mcp_server import operations as mcp_operations
"""Unit tests for MCP OAuth passthrough tool-fetch behavior."""
import asyncio
import logging
import sys
from unittest.mock import AsyncMock, MagicMock
@ -462,7 +464,6 @@ async def test_fetch_tools_logs_upstream_request_details_on_500(caplog):
assert "upstream-token-0123456789" not in caplog.text
@pytest.mark.asyncio
async def test_client_creation_failure_logs_sanitized_exchange(monkeypatch, caplog):
manager = MCPServerManager()
@ -498,3 +499,243 @@ def test_passthrough_admission_recognizes_only_matching_authorization(oauth_head
server = MCPServer(server_id="catalog", name="catalog", alias="catalog", transport=MCPTransport.http)
assert _client_has_passthrough_authorization(server, oauth_headers, server_headers) is authorized
@pytest.mark.parametrize("base_url, request_path, scope, expected", [
("https://gateway.test/proxy", "/alias/mcp", "tools.write", 'Bearer resource_metadata="https://gateway.test/proxy/.well-known/oauth-protected-resource/alias/mcp", error="insufficient_scope", scope="tools.write"'),
("https://gateway.test/proxy", "/mcp/alias", "", 'Bearer resource_metadata="https://gateway.test/proxy/.well-known/oauth-protected-resource/mcp/alias", error="insufficient_scope"'),
(None, None, "tools.write", 'Bearer error="insufficient_scope", scope="tools.write"'),
])
def test_managed_step_up_challenge_uses_only_the_gateway_metadata(base_url, request_path, scope, expected):
error = MCPUpstreamAuthError(status_code=403, www_authenticate=None, server_name="alias", required_scope=scope)
response = error.to_http_exception(base_url=base_url, request_path=request_path)
assert response.status_code == 403
assert response.headers == {"www-authenticate": expected}
@pytest.mark.asyncio
async def test_scope_challenge_replaces_pending_stream_headers_once():
from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse
send = AsyncMock()
response = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp")
await response({"type": "http.response.start", "status": 200, "headers": []})
body = asyncio.create_task(response({"type": "http.response.body", "body": b"event: message\n"}))
await asyncio.sleep(0)
send.assert_not_awaited()
response.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write"))
await asyncio.wait_for(body, 1)
await response({"type": "http.response.body", "body": b"private tool error"})
assert send.await_count == 2
start, result = [call.args[0] for call in send.await_args_list]
assert start["status"] == 403
challenge = dict(start["headers"])[b"www-authenticate"].decode()
assert challenge == ('Bearer resource_metadata="https://gateway.example/.well-known/oauth-protected-resource/docs/mcp", '
'error="insufficient_scope", scope="tools.write"')
assert result["body"] == b'{"detail":"Forbidden"}'
@pytest.mark.asyncio
async def test_scope_authorization_releases_progress_without_waiting_for_tool_result():
import httpx2
from mcp.server.context import ServerRequestContext
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
from litellm.proxy._experimental.mcp_server.scope_challenge import (
SCOPE_RESPONSE_KEY,
OAuthScopeResponse,
)
from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response, record_upstream_tool_authorization
send = AsyncMock()
gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs")
request = Request({"type": "http", SCOPE_RESPONSE_KEY: gate})
context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=request)
token = active_mcp_request_ctx_var.set(context)
try:
await gate({"type": "http.response.start", "status": 200, "headers": []})
progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True}))
await asyncio.sleep(0)
for method, status in [("initialize", 200), ("tools/list", 200), ("tools/call", 403)]:
await record_upstream_tool_authorization(httpx2.Response(status, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": method})), pending=get_scope_response())
assert gate.pending
unread = httpx2.Request("POST", "https://upstream.example/mcp", content=iter([b'{"method":"tools/call"}']))
await record_upstream_tool_authorization(httpx2.Response(200, request=unread), pending=get_scope_response())
assert gate.pending
# Use the actual HTTP response hook contract, before reading the tool response body.
from litellm.experimental_mcp_client.client import MCPClient
upstream = MCPClient(server_url="https://upstream.example/mcp")
async with upstream._create_httpx_client_factory(transport=httpx2.MockTransport(lambda request: httpx2.Response(200, content=b"tool result")))() as client:
async with client.stream("POST", "https://upstream.example/mcp", json={"method": "tools/call"}):
await asyncio.wait_for(progress, 1)
assert send.await_args_list[-1].args[0]["body"] == b"progress"
gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write"))
await gate({"type": "http.response.body", "body": b"result"})
assert send.await_args_list[0].args[0]["status"] == 200
assert send.await_args_list[-1].args[0]["body"] == b"result"
finally:
active_mcp_request_ctx_var.reset(token)
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError])
async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation(failure):
from mcp.server.context import ServerRequestContext
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response
send = AsyncMock()
gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs")
context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}))
with pytest.raises(failure):
await finish_scope_response(context, AsyncMock(side_effect=failure))
await gate({"type": "http.response.start", "status": 200, "headers": []})
await gate({"type": "http.response.body", "body": b"error"})
assert not gate.pending
assert send.await_args_list[-1].args[0]["body"] == b"error"
@pytest.mark.asyncio
async def test_scope_response_preserves_transport_validation_errors():
from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse
send = AsyncMock()
gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs")
await gate({"type": "http.response.start", "status": 400, "headers": []})
await gate({"type": "http.response.body", "body": b"invalid request"})
assert send.await_args_list[0].args[0]["status"] == 400
assert send.await_args_list[1].args[0]["body"] == b"invalid request"
@pytest.mark.asyncio
@pytest.mark.parametrize("denied", [False, True])
@pytest.mark.parametrize("stateless", [False, True])
@pytest.mark.parametrize("padding_size", [0, 131072])
async def test_streamable_sdk_relays_scope_challenge_before_committing_http_response(denied, stateless, padding_size):
import httpx2
from mcp.server.lowlevel import Server
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.types import CallToolResult, TextContent
from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response
from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response
async def call_tool(context, params):
token = active_mcp_request_ctx_var.set(context)
try:
gate = get_scope_response()
assert gate is not None
if denied:
gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write"))
return CallToolResult(content=[TextContent(type="text", text="done")])
finally:
active_mcp_request_ctx_var.reset(token)
server = Server("scope-test", on_call_tool=call_tool)
server.middleware.append(finish_scope_response)
manager = StreamableHTTPSessionManager(server, stateless=stateless)
async def app(scope, receive, send):
response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs")
scope[SCOPE_RESPONSE_KEY] = response
await manager.handle_request(scope, receive, response)
async with manager.run():
async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client:
headers = {"Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2025-11-25"}
if not stateless:
initialized = await asyncio.wait_for(client.post("https://gateway.example/mcp/docs", headers=headers, json={
"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {
"protocolVersion": "2025-11-25", "capabilities": {}, "clientInfo": {"name": "scope-test" + "x" * padding_size, "version": "1"},
},
}), 3)
assert initialized.status_code == 200
headers["MCP-Session-Id"] = initialized.headers["mcp-session-id"]
notified = await client.post("https://gateway.example/mcp/docs", headers=headers, json={"jsonrpc": "2.0", "method": "notifications/initialized"})
assert notified.status_code == 202
response = await asyncio.wait_for(client.post("https://gateway.example/mcp/docs", headers=headers,
json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "write", "arguments": {"padding": "x" * padding_size}}}), 3)
assert response.status_code == (403 if denied else 200)
if denied:
assert 'error="insufficient_scope"' in response.headers["www-authenticate"]
assert response.json() == {"detail": "Forbidden"}
else:
assert '"text":"done"' in response.text
assert "www-authenticate" not in response.headers
@pytest.mark.asyncio
async def test_scope_challenges_are_isolated_between_concurrent_requests():
import httpx2
from mcp.server.context import ServerRequestContext
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server.legacy_callbacks import get_scope_response, record_upstream_tool_authorization
from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse
async def request_scope(denied):
send = AsyncMock()
gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs")
context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}))
token = active_mcp_request_ctx_var.set(context)
try:
await gate({"type": "http.response.start", "status": 200, "headers": []})
await asyncio.sleep(0)
assert get_scope_response() is gate
if denied:
gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write"))
else:
await record_upstream_tool_authorization(httpx2.Response(200, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": "tools/call"})), pending=get_scope_response())
await gate({"type": "http.response.body", "body": b"done"})
return send.await_args_list[0].args[0]["status"]
finally:
active_mcp_request_ctx_var.reset(token)
assert await asyncio.gather(request_scope(True), request_scope(False)) == [403, 200]
assert get_scope_response() is None
@pytest.mark.parametrize("body", [b"", b"{", b"[]", b'{"method":"initialize"}'])
def test_non_tool_payloads_do_not_hold_response_headers(body):
from litellm.proxy._experimental.mcp_server.scope_challenge import is_tool_call_request, scope_response_from_request
assert not is_tool_call_request(body)
assert scope_response_from_request(None) is None
@pytest.mark.asyncio
async def test_scope_middleware_preserves_non_http_requests():
from mcp.server.context import ServerRequestContext
from mcp.types import CallToolResult, TextContent
from litellm.proxy._experimental.mcp_server.scope_challenge import finish_scope_response
context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call")
async def call_next(request):
return CallToolResult(content=[TextContent(type="text", text=request.method)])
returned = await finish_scope_response(context, call_next)
assert returned.content[0].text == "tools/call"
@pytest.mark.parametrize("path", ["/mcp", "/toolset/demo/mcp"])
def test_managed_step_up_metadata_matches_the_sdk_configured_resource(path):
from mcp.shared.auth_utils import check_resource_allowed
from urllib.parse import urlsplit
error = MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")
response = error.to_http_exception("https://gateway.test", path)
challenge = response.headers["www-authenticate"]
metadata = challenge.split('resource_metadata="', 1)[1].split('"', 1)[0]
metadata_path = urlsplit(metadata).path
resource_path = metadata_path.removeprefix("/.well-known/oauth-protected-resource")
assert check_resource_allowed(
requested_resource=f"https://gateway.test{path}",
configured_resource=f"https://gateway.test{resource_path}",
), "the SDK must accept the advertised resource before it can reauthorize"

83
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -2211,6 +2211,31 @@ class TestMCPServerManager:
assert exc_info.value.www_authenticate == challenge
assert mock_client.call_tool.call_args.kwargs.get("raise_on_error") is True
@pytest.mark.asyncio
async def test_managed_scope_challenge_requires_reauthorization_without_retry(self):
server = MCPServer(
server_id="managed-scope", name="managed-scope", alias="managed-scope",
url="https://up.example.com/mcp", transport=MCPTransport.http,
auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code",
)
upstream_error = self._upstream_status_error(
403, 'Bearer error="insufficient_scope", scope="tools.write", resource_metadata="https://untrusted.test/metadata"'
)
manager = MCPServerManager()
client = AsyncMock()
client.call_tool = AsyncMock(side_effect=upstream_error)
manager.create_mcp_client = AsyncMock(return_value=client)
with pytest.raises(MCPUpstreamAuthError) as raised:
await self._run_call_regular(manager, server)
response = raised.value.to_http_exception("https://gateway.test", "/managed-scope/mcp")
assert response.status_code == 403
assert response.headers == {
"www-authenticate": 'Bearer resource_metadata="https://gateway.test/.well-known/oauth-protected-resource/managed-scope/mcp", error="insufficient_scope", scope="tools.write"'
}
assert client.call_tool.await_count == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("is_error", [False, True])
async def test_call_passthrough_returns_tool_result_unchanged(self, is_error):
@ -4533,6 +4558,7 @@ class TestMCPServerManager:
f"{issuer}/.well-known/oauth-authorization-server",
{
"issuer": issuer,
"response_types_supported": ["code"],
"authorization_endpoint": "https://idp.example.com/authorize",
"token_endpoint": "https://idp.example.com/token",
"scopes_supported": ["read", "write"],
@ -4554,6 +4580,11 @@ class TestMCPServerManager:
assert result.client_id_metadata_document_supported is True
assert result.authorization_server_metadata is not None
assert result.authorization_server_metadata["scopes_supported"] == ["read", "write"]
merged = manager._merge_discovered_oauth_metadata(self._passthrough_call_server(MCPAuth.oauth2), result)
assert merged.authorization_server_metadata == result.authorization_server_metadata
@pytest.mark.asyncio
async def test_fetch_single_authorization_server_metadata_rejects_issuer_mismatch(self):
"""RFC 8414 §3.3 fail-closed: a document self-attesting a DIFFERENT issuer than the one it was
@ -7255,6 +7286,47 @@ class TestMCPServerManager:
def _unrestricted_auth() -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key="sk-test")
@pytest.mark.asyncio
@pytest.mark.parametrize("managed", [False, True])
async def test_tool_response_waits_only_for_managed_user_oauth_authorization(self, managed):
import asyncio
from mcp.server.context import ServerRequestContext
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse
auth = self._unrestricted_auth()
manager, logging = self._manager_ready_for_call_tool(
[MCPTool(name="test_tool", inputSchema={"type": "object"})],
caller=ListedToolsCaller(user_api_key_auth=auth),
)
if managed:
manager.registry["test-server"] = manager.registry["test-server"].model_copy(update={
"auth_type": MCPAuth.oauth2, "oauth2_flow": "authorization_code",
})
send = AsyncMock()
gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/test-server")
context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}))
token = active_mcp_request_ctx_var.set(context)
await gate({"type": "http.response.start", "status": 200, "headers": []})
progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True}))
async def upstream_call(*args, **kwargs):
await asyncio.sleep(0)
assert send.await_count == (0 if managed else 2)
return CallToolResult(content=[], is_error=False)
manager.create_mcp_client.return_value.call_tool.side_effect = upstream_call
try:
result = await manager.call_tool(server_name="test-server", name="test_tool", arguments={}, user_api_key_auth=auth, proxy_logging_obj=logging)
assert not result.is_error
finally:
gate.allow()
await asyncio.wait_for(progress, 1)
active_mcp_request_ctx_var.reset(token)
@pytest.mark.asyncio
async def test_call_tool_hands_listed_tool_description_and_schema_to_pre_call_hooks(self):
schema = {"type": "object", "properties": {"param": {"type": "string"}}, "required": ["param"]}
@ -18051,7 +18123,10 @@ async def test_unreusable_upstream_session_is_replaced_before_cleanup_finishes(
@pytest.fixture
def _active_mcp_request_context() -> Iterator[None]:
token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object()))
context: Final = ServerRequestContext(
session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=None
)
token: Final = active_mcp_request_ctx_var.set(context)
try:
yield
finally:
@ -19973,17 +20048,23 @@ def test_oauth_rebuild_retains_only_corroborated_cimd_capability(capability: boo
transport=MCPTransport.http, auth_type=MCPAuth.oauth2,
authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token",
client_id_metadata_document_supported=capability,
authorization_server_metadata={"issuer": "https://idp.example.com", "response_types_supported": ["code"], "scopes_supported": ["offline_access"]},
)
rebuilt: Final = previous.model_copy(update={
"client_id_metadata_document_supported": not capability if rebuild == "fresh_discovery" else None,
"authorization_url": "https://changed.example.com/authorize" if rebuild == "repointed" else previous.authorization_url,
"token_url": None,
"authorization_server_metadata": {"scopes_supported": []} if rebuild == "fresh_discovery" else None,
"issuer": "https://idp.example.com" if rebuild == "anchored" else None,
"issuer_is_anchored": rebuild == "anchored",
})
carry_forward_resolved_oauth_endpoints(rebuilt, previous)
expected: Final = not capability if rebuild == "fresh_discovery" else None if rebuild in ("repointed", "anchored") else capability
assert rebuilt.client_id_metadata_document_supported is expected
assert rebuilt.authorization_server_metadata == (
{"scopes_supported": []} if rebuild == "fresh_discovery"
else previous.authorization_server_metadata if rebuild == "same" else None
)
@pytest.mark.asyncio

45
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
View file

@ -264,11 +264,19 @@ async def test_mcp_server_tool_call_strips_custom_litellm_key_header(_mcp_reques
@pytest.mark.asyncio
async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_request_ctx):
"""The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session
tool call cannot emit a raw 401 the way the REST path does. mcp_server_tool_call must turn an
upstream MCPUpstreamAuthError into an explicit isError result naming the status, not a masked
500 or a raw traceback, so the client still learns it must re-authenticate upstream."""
@pytest.mark.parametrize("required_scope, http_transport", [(None, False), (None, True), ("tools.write", True)])
async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(
_mcp_request_ctx, required_scope, http_transport
):
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse
status_code = 403 if required_scope is not None else 401
send = AsyncMock()
response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/pt")
request = Request({"type": "http", "headers": [], SCOPE_RESPONSE_KEY: response})
await response({"type": "http.response.start", "status": 200, "headers": []})
try:
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
from litellm.proxy._experimental.mcp_server.server import (
@ -284,7 +292,9 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r
return data
async def mock_call_mcp_tool(*args, **kwargs):
raise MCPUpstreamAuthError(status_code=401, www_authenticate="Bearer", server_name="pt")
raise MCPUpstreamAuthError(
status_code=status_code, www_authenticate=None, server_name="pt", required_scope=required_scope
)
mock_logger = MagicMock()
with patch(
@ -298,20 +308,27 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r
with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()):
with patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger", mock_logger):
result = await mcp_server_tool_call(
_mcp_request_ctx(), _call_tool_params("test_tool", {"param": "value"})
_mcp_request_ctx(request=request if http_transport else None),
_call_tool_params("test_tool", {"param": "value"}),
)
assert result.is_error is True
# The dedicated MCPUpstreamAuthError branch (not the generic Exception fallthrough) produces this
# specific message and logs at info, never a traceback via verbose_logger.exception.
assert "upstream authentication required" in result.content[0].text
assert "401" in result.content[0].text
assert str(status_code) in result.content[0].text
exception_calls = [str(c.args[0]) for c in mock_logger.exception.call_args_list if c.args]
assert not any("mcp_server_tool_call" in m for m in exception_calls), (
"must not log a traceback for the expected re-auth"
)
info_calls = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args]
assert any("Upstream auth failure" in m for m in info_calls)
response.allow()
await response({"type": "http.response.body", "body": b"tool error"})
start = send.await_args_list[0].args[0]
assert start["status"] == (403 if required_scope is not None else 200)
if required_scope is not None:
assert b'error="insufficient_scope"' in dict(start["headers"])[b"www-authenticate"]
def test_prepare_mcp_server_headers_case_insensitive_extra_headers():
@ -2464,9 +2481,12 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body():
(("é", 0), ("é", 1), ("中", 1), ("中", 2), ("😀", 1), ("😀", 2), ("😀", 3)),
)
async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap(
method: str, chunked: bool, character: str, bytes_before_cap: int
method: str, chunked: bool, character: str, bytes_before_cap: int, _mcp_request_ctx
) -> None:
from starlette.requests import Request
from litellm.proxy._experimental.mcp_server import server as mcp_module
from litellm.proxy._experimental.mcp_server.scope_challenge import finish_scope_response
params: Final = (
{
@ -2490,9 +2510,14 @@ async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap(
send: Final = AsyncMock()
received: Final[asyncio.Future[bytes]] = asyncio.get_running_loop().create_future()
async def handle_request(_: Scope, downstream_receive: Receive, outgoing: Send) -> None:
async def handle_request(incoming: Scope, downstream_receive: Receive, outgoing: Send) -> None:
assert receive.await_count == (2 if chunked else 1)
received.set_result(await _drain_body(downstream_receive))
async def dispatch(context):
return CallToolResult(content=[])
await finish_scope_response(_mcp_request_ctx(request=Request(incoming)), dispatch)
await outgoing({"type": "http.response.start", "status": 200, "headers": []})
await outgoing({"type": "http.response.body", "body": b"{}"})

91
ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.integration.test.tsx Normal file
View file

@ -0,0 +1,91 @@
import { webcrypto } from "node:crypto";
import { afterEach, expect, it, vi } from "vitest";
import { render, screen, waitFor } from "@testing-library/react";
import ConnectFlowBanner from "./ConnectFlowBanner";
const locationDescriptor = Object.getOwnPropertyDescriptor(window, "location")!;
afterEach(() => {
vi.unstubAllGlobals();
window.sessionStorage.clear();
Object.defineProperty(window, "location", locationDescriptor);
});
it("carries a challenged scope through the connect button into vendor authorization", async () => {
vi.stubGlobal("crypto", webcrypto);
vi.stubGlobal(
"fetch",
vi.fn(async () => Response.json({ client_id: "registered-client" })),
);
Object.defineProperty(window, "location", {
configurable: true,
value: {
href: "https://gateway.example/ui/connect?connect_flow=flow-handle",
origin: "https://gateway.example",
pathname: "/ui/connect",
},
});
render(
<ConnectFlowBanner
flowHandle="flow-handle"
flow={{
state: "interactive",
client_origin: "https://client.example",
server_id: "server-1",
server_name: "Scoped tools",
connected: false,
requested_scopes: ["tools.write"],
}}
accessToken="caller"
onConnected={vi.fn()}
failed={false}
/>,
);
await waitFor(() => expect(window.location.href).toContain("scope=tools.write"));
expect(new URL(window.location.href, window.location.origin).pathname).toBe(
"/v1/mcp/server/oauth/server-1/authorize",
);
expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument();
});
it("starts consent for the next upstream in the same connect flow", async () => {
vi.stubGlobal("crypto", webcrypto);
vi.stubGlobal(
"fetch",
vi.fn(async () => Response.json({ client_id: "registered-client" })),
);
Object.defineProperty(window, "location", {
configurable: true,
value: {
href: "https://gateway.example/ui/connect?connect_flow=flow-handle",
origin: "https://gateway.example",
pathname: "/ui/connect",
},
});
const props = {
flowHandle: "flow-handle",
accessToken: "caller",
onConnected: vi.fn(),
failed: false,
};
const first = {
state: "interactive" as const,
client_origin: "https://client.example",
server_id: "first",
server_name: "First",
connected: false,
requested_scopes: ["first.write"],
};
const { rerender } = render(<ConnectFlowBanner {...props} flow={first} />);
await waitFor(() => expect(window.location.href).toContain("scope=first.write"));
rerender(
<ConnectFlowBanner
{...props}
flow={{ ...first, server_id: "second", server_name: "Second", requested_scopes: ["second.write"] }}
/>,
);
await waitFor(() => expect(window.location.href).toContain("scope=second.write"));
expect(new URL(window.location.href, window.location.origin).pathname).toBe("/v1/mcp/server/oauth/second/authorize");
expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
});

4
ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx
View file

@ -76,11 +76,13 @@ const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, flow, accessToken, onC
<div className="flex shrink-0 gap-2">
{vendorServer !== null && (
<OAuth2ConnectButton
key={vendorServer.server_id}
server={vendorServer}
accessToken={accessToken}
onConnect={onConnected}
scopes={flow?.requested_scopes}
variant="button"
autoStartKey={`litellm-mcp-autostart:${flowHandle}`}
autoStartKey={`litellm-mcp-autostart:${flowHandle}:${vendorServer.server_id}`}
/>
)}
<form method="POST" action={action}>

8
ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx
View file

@ -31,6 +31,7 @@ interface OAuth2ConnectButtonProps {
onConnect: (serverId: string) => void;
variant?: "badge" | "button";
autoStartKey?: string | null;
scopes?: string[];
}
export const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
@ -39,14 +40,17 @@ export const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
onConnect,
variant = "badge",
autoStartKey = null,
scopes,
}) => {
const name = server.server_name ?? server.alias ?? server.server_id;
const { startOAuthFlow, status } = useUserMcpOAuthFlow({
const flowOptions = {
accessToken,
serverId: server.server_id,
serverAlias: name,
scopes,
onSuccess: useCallback(() => onConnect(server.server_id), [onConnect, server.server_id]),
});
};
const { startOAuthFlow, status } = useUserMcpOAuthFlow(flowOptions);
useEffect(() => {
if (autoStartKey === null || status !== "idle" || getSecureItem(autoStartKey) !== null) return;

1
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -4782,6 +4782,7 @@ export interface ConnectFlowStatus {
server_id: string | null;
server_name: string | null;
connected: boolean | null;
requested_scopes?: string[];
}
/**

71
ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.test.tsx Normal file
View file

@ -0,0 +1,71 @@
import { renderHook, waitFor } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import * as networking from "@/components/networking";
import { setSecureItem } from "@/utils/secureStorage";
import { useUserMcpOAuthFlow } from "./useUserMcpOAuthFlow";
vi.mock("@/components/networking", () => ({
exchangeMcpOAuthToken: vi.fn(),
storeMCPOAuthUserCredential: vi.fn(),
registerMcpOAuthClient: vi.fn(),
buildMcpOAuthAuthorizeUrl: vi.fn(),
getProxyBaseUrl: () => "",
serverRootPath: "",
}));
vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), error: vi.fn() } }));
describe("user OAuth granted scopes", () => {
beforeEach(() => {
vi.clearAllMocks();
window.sessionStorage.clear();
const flowState = {
state: "callback-state",
codeVerifier: "verifier",
serverId: "server-1",
redirectUri: "https://app.example/callback",
scopes: ["tools.write"],
};
setSecureItem("litellm-user-mcp-oauth-flow-state", JSON.stringify(flowState));
setSecureItem(
"litellm-user-mcp-oauth-result",
JSON.stringify({
state: "callback-state",
code: "provider-code",
}),
);
});
it.each([
["tools.read tools.write", ["tools.read", "tools.write"]],
["tools.read", ["tools.read"]],
["", []],
[undefined, ["tools.write"]],
])("persists the actual grant %s after reconnect", async (scope, expected) => {
vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue({
access_token: "vendor-token",
refresh_token: "vendor-refresh",
scope,
});
const onSuccess = vi.fn();
const { result } = renderHook(() =>
useUserMcpOAuthFlow({
accessToken: "caller",
serverId: "server-1",
onSuccess,
}),
);
await waitFor(() => expect(result.current.status).toBe("success"));
const expectedCredential = {
access_token: "vendor-token",
refresh_token: "vendor-refresh",
expires_in: undefined,
scopes: expected,
};
expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalledExactlyOnceWith(
"caller",
"server-1",
expectedCredential,
);
expect(onSuccess).toHaveBeenCalledOnce();
});
});

5
ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx
View file

@ -68,6 +68,9 @@ const getStorage = (key: string): string | null => {
return getSecureItem(key);
};
const grantedScopes = (scope: unknown, requested: string[] | undefined): string[] | undefined =>
typeof scope === "string" ? scope.split(" ").filter(Boolean) : requested;
export const useUserMcpOAuthFlow = ({
accessToken,
serverId,
@ -216,7 +219,7 @@ export const useUserMcpOAuthFlow = ({
access_token: token.access_token,
refresh_token: token.refresh_token,
expires_in: token.expires_in,
scopes: flowState.scopes,
scopes: grantedScopes(token.scope, flowState.scopes),
});
setStatus("success");

1
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -51945,6 +51945,7 @@ export interface operations {
code_challenge_method?: string | null;
response_type?: string | null;
resource?: string | null;
scope?: string | null;
};
header?: never;
path?: never;