scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge pull request #40791 from BerriAI/litellm_fix_mcp_root_discovery_6634

Browse source 
fix(mcp): use gateway authentication for root discovery
This commit is contained in:
joshua-berri 2026-09-11 20:23:10 -07:00 • committed by GitHub
commit 70cf348aa5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 585 additions and 188 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py
View file

@ -7,7 +7,6 @@ just a form that asks the user for their API key — not a full identity-provide
Endpoints implemented here:
GET /.well-known/oauth-authorization-server — OAuth authorization server metadata
GET /.well-known/oauth-protected-resource — OAuth protected resource metadata
GET /v1/mcp/oauth/authorize — Shows HTML form to collect the API key
POST /v1/mcp/oauth/authorize — Stores temp auth code and redirects
POST /v1/mcp/oauth/token — Exchanges code for a bearer JWT token
@ -19,7 +18,7 @@ import html as _html_module
import time
import uuid
from typing import Final, cast
from urllib.parse import urlencode
from urllib.parse import urlencode, urlparse
import jwt
from fastapi import APIRouter, Depends, Form, HTTPException, Request
@ -27,14 +26,15 @@ from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
from litellm._logging import verbose_proxy_logger
from litellm.proxy._experimental.mcp_server.db import store_user_credential
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
get_request_base_url,
)
from litellm.proxy._experimental.mcp_server.oauth_utils import (
BYOK_RESOURCE_METADATA_PATH,
TOKEN_NO_CACHE_HEADERS,
get_request_base_url,
validate_loopback_redirect_uri,
well_known_root_suffix,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.middleware.per_request_root_path_middleware import get_server_root_paths
# ---------------------------------------------------------------------------
# In-memory store for pending authorization codes.
@ -596,13 +596,10 @@ def _build_authorize_html(
# ---------------------------------------------------------------------------
@router.get("/.well-known/oauth-authorization-server", include_in_schema=False)
async def oauth_authorization_server_metadata(request: Request) -> JSONResponse:
"""RFC 8414 Authorization Server Metadata for the BYOK OAuth flow."""
base_url: Final = get_request_base_url(request)
def _byok_authorization_server_response(base_url: str, issuer: str) -> JSONResponse:
return JSONResponse(
{
"issuer": base_url,
"issuer": issuer,
"authorization_endpoint": f"{base_url}/v1/mcp/oauth/authorize",
"token_endpoint": f"{base_url}/v1/mcp/oauth/token",
"response_types_supported": ["code"],
@ -612,14 +609,36 @@ async def oauth_authorization_server_metadata(request: Request) -> JSONResponse:
)
@router.get("/.well-known/oauth-protected-resource", include_in_schema=False)
async def oauth_protected_resource_metadata(request: Request) -> JSONResponse:
"""RFC 9728 Protected Resource Metadata pointing back at this server."""
@router.get("/.well-known/oauth-authorization-server", include_in_schema=False)
async def oauth_authorization_server_metadata(request: Request) -> JSONResponse:
base_url: Final = get_request_base_url(request)
return _byok_authorization_server_response(base_url, base_url)
@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/v1/mcp/oauth", include_in_schema=False)
async def byok_authorization_server_metadata(request: Request) -> JSONResponse:
base_url: Final = get_request_base_url(request)
return _byok_authorization_server_response(base_url, f"{base_url}/v1/mcp/oauth")
@router.get("/.well-known/oauth-authorization-server/{root_path:path}/v1/mcp/oauth", include_in_schema=False)
async def byok_prefixed_authorization_server_metadata(request: Request, root_path: str) -> JSONResponse:
prefix: Final = f"/{root_path}"
if prefix not in get_server_root_paths():
raise HTTPException(status_code=404, detail="Unknown proxy root path")
parsed: Final = urlparse(get_request_base_url(request))
base_url: Final = f"{parsed.scheme}://{parsed.netloc}{prefix}"
return _byok_authorization_server_response(base_url, f"{base_url}/v1/mcp/oauth")
@router.get(BYOK_RESOURCE_METADATA_PATH, include_in_schema=False)
async def byok_protected_resource_metadata(request: Request) -> JSONResponse:
base_url: Final = get_request_base_url(request)
parsed: Final = urlparse(base_url)
return JSONResponse(
{
"resource": base_url,
"authorization_servers": [base_url],
"resource": f"{parsed.scheme}://{parsed.netloc}",
"authorization_servers": (f"{base_url}/v1/mcp/oauth",),
}
)

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

@ -1842,6 +1842,30 @@ async def register_client_with_server(
return JSONResponse(token_response)
@router.get("/authorize/mcp-session")
async def authorize_mcp_session(
request: Request,
redirect_uri: str,
client_id: str,
state: str = "",
code_challenge: str | None = None,
code_challenge_method: str | None = None,
response_type: str | None = None,
resource: str | None = None,
) -> Response:
return aggregate_authorize(
request=request,
client_id=client_id,
redirect_uri=redirect_uri,
state=state,
code_challenge=code_challenge,
code_challenge_method=code_challenge_method,
response_type=response_type,
session_user_id=_session_cookie_user_id(request),
resource=resource,
)
@router.get("/{mcp_server_name}/authorize")
@router.get("/authorize")
async def authorize(
@ -2393,8 +2417,7 @@ async def _build_oauth_protected_resource_response(
per-server URL completes the same sign-in flow the aggregate ``/mcp`` endpoint
supports and is admitted with a gateway session bearer. The per-server relay
authorize/token endpoints stay registered for the keyed interactive flow (which
is challenged with an explicit ``authorization_uri``), and the root-resolved
(unnamed) legacy shape keeps the relay authorization server.
is challenged with an explicit ``authorization_uri``).
Args:
request: FastAPI Request object
@ -2405,15 +2428,11 @@ async def _build_oauth_protected_resource_response(
Returns:
OAuth protected resource metadata dict
"""
if mcp_server_name is None:
return oauth_protected_resource_root(request)
request_base_url: Final = get_request_base_url(request)
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
explicitly_named: Final = mcp_server_name is not None
# When no server name provided, try to resolve the single OAuth2 server
if mcp_server_name is None:
resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if resolved:
mcp_server_name = resolved.server_name or resolved.name
mcp_server: MCPServer | None = None
if mcp_server_name:
@ -2478,7 +2497,7 @@ async def _build_oauth_protected_resource_response(
if obo_response is not None:
return obo_response
if explicitly_named and mcp_server is not None and mcp_server.advertises_gateway_authorization_server:
if mcp_server is not None and mcp_server.advertises_gateway_authorization_server:
return {
"authorization_servers": [f"{request_base_url}/mcp"],
"resource": resource_url,
@ -2542,6 +2561,17 @@ def _jwt_auth_issuers() -> list:
return issuers
@router.get("/.well-known/oauth-protected-resource")
def oauth_protected_resource_root(request: Request) -> dict[str, str | tuple[str, ...]]:
request_base_url: Final = get_request_base_url(request)
parsed: Final = urlparse(request_base_url)
return {
"resource": f"{parsed.scheme}://{parsed.netloc}",
"authorization_servers": (f"{request_base_url}/mcp",),
"scopes_supported": (),
}
def _build_aggregate_protected_resource_response(request: Request) -> dict:
"""RFC 9728 metadata for the aggregate /mcp resource: the gateway itself is
the authorization server. No per-server names or scopes leak here; access
@ -2568,14 +2598,14 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict:
The issuer is ``{base}/mcp`` and must stay equal to the value the
aggregate protected-resource document advertises: spec clients verify the
issuer in the metadata matches the one that derived the well-known URL.
Advertises the root /authorize, /token, and /register endpoints and
Advertises the MCP session authorize endpoint, root /token and /register endpoints, and
``token_endpoint_auth_methods_supported: ["none", ...]`` because DCR
clients (Claude Desktop, MCP Inspector) register as public clients; PKCE
S256 is mandatory in the gateway's authorize flow."""
request_base_url: Final = get_request_base_url(request)
return {
"issuer": f"{request_base_url}/mcp",
"authorization_endpoint": f"{request_base_url}/authorize",
"authorization_endpoint": f"{request_base_url}/authorize/mcp-session",
"token_endpoint": f"{request_base_url}/token",
"introspection_endpoint": f"{request_base_url}/introspect",
"registration_endpoint": f"{request_base_url}/register",
@ -2645,7 +2675,6 @@ async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_nam
# LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp
# Kept for backward compatibility with existing deployments
@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/{{mcp_server_name}}/mcp")
@router.get("/.well-known/oauth-protected-resource")
async def oauth_protected_resource_mcp(request: Request, mcp_server_name: str | None = None):
"""
OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.

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

@ -94,6 +94,7 @@ from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
from litellm.proxy._experimental.mcp_server.oauth_utils import (
_redact_mcp_resource_url,
canonicalize_url_identity,
get_byok_www_authenticate,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials import (
Error,
@ -1234,7 +1235,7 @@ async def _resolve_byok_mcp_auth_header(
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'},
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
return byok_cred

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

@ -16,6 +16,7 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
normalize_token_endpoint_auth_method,
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.middleware.per_request_root_path_middleware import get_request_root_path
if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -126,6 +127,14 @@ def _resolve_proxy_base_url_env() -> str | None:
return None
BYOK_RESOURCE_METADATA_PATH: Final = "/v1/mcp/oauth/protected-resource"
def get_byok_www_authenticate() -> str:
base_url: Final = _resolve_proxy_base_url_env() or get_request_root_path().rstrip("/")
return f'Bearer resource_metadata="{base_url}{BYOK_RESOURCE_METADATA_PATH}"'
def get_request_base_url(request: Request) -> str:
"""
Get the base URL for the request, considering X-Forwarded-* headers.

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

@ -56,6 +56,7 @@ from litellm.proxy._experimental.mcp_server.mcp_debug import (
)
from litellm.proxy._experimental.mcp_server.oauth_utils import (
_redact_mcp_resource_url,
get_byok_www_authenticate,
get_passthrough_www_authenticate,
get_route_relative_request_path,
well_known_root_suffix,
@ -2852,7 +2853,7 @@ if MCP_AVAILABLE:
"server_name": mcp_server.server_name or mcp_server.name,
"message": "User identity is required for BYOK servers",
},
headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'},
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
# Check shared credential cache before hitting the DB.
@ -2873,9 +2874,7 @@ if MCP_AVAILABLE:
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={
"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'
},
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
return
@ -2914,7 +2913,7 @@ if MCP_AVAILABLE:
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'},
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
async def execute_mcp_tool(
@ -3068,9 +3067,7 @@ if MCP_AVAILABLE:
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={
"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'
},
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
mcp_auth_header = byok_cred
elif mcp_server.is_byok:

221
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -22415,47 +22415,34 @@
},
"/.well-known/oauth-protected-resource": {
"get": {
"description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.",
"operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource_get",
"parameters": [
{
"in": "query",
"name": "mcp_server_name",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Mcp Server Name"
}
}
],
"operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
"additionalProperties": {
"anyOf": [
{
"type": "string"
},
{
"items": {
"type": "string"
},
"type": "array"
}
]
},
"title": "Response Oauth Protected Resource Root Well Known Oauth Protected Resource Get",
"type": "object"
}
}
},
"description": "Validation Error"
"description": "Successful Response"
}
},
"summary": "Oauth Protected Resource Mcp",
"summary": "Oauth Protected Resource Root",
"tags": [
"mcp_byok_oauth"
]
@ -24636,47 +24623,34 @@
},
"/.well-known/oauth-protected-resource": {
"get": {
"description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.",
"operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource_get_2",
"parameters": [
{
"in": "query",
"name": "mcp_server_name",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Mcp Server Name"
}
}
],
"operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2",
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
"additionalProperties": {
"anyOf": [
{
"type": "string"
},
{
"items": {
"type": "string"
},
"type": "array"
}
]
},
"title": "Response Oauth Protected Resource Root Well Known Oauth Protected Resource Get",
"type": "object"
}
}
},
"description": "Validation Error"
"description": "Successful Response"
}
},
"summary": "Oauth Protected Resource Mcp",
"summary": "Oauth Protected Resource Root",
"tags": [
"mcp_discoverable"
]
@ -25052,6 +25026,129 @@
]
}
},
"/authorize/mcp-session": {
"get": {
"operationId": "authorize_mcp_session_authorize_mcp_session_get",
"parameters": [
{
"in": "query",
"name": "redirect_uri",
"required": true,
"schema": {
"title": "Redirect Uri",
"type": "string"
}
},
{
"in": "query",
"name": "client_id",
"required": true,
"schema": {
"title": "Client Id",
"type": "string"
}
},
{
"in": "query",
"name": "state",
"required": false,
"schema": {
"default": "",
"title": "State",
"type": "string"
}
},
{
"in": "query",
"name": "code_challenge",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Code Challenge"
}
},
{
"in": "query",
"name": "code_challenge_method",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Code Challenge Method"
}
},
{
"in": "query",
"name": "response_type",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Response Type"
}
},
{
"in": "query",
"name": "resource",
"required": false,
"schema": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Resource"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"summary": "Authorize Mcp Session",
"tags": [
"mcp_discoverable"
]
}
},
"/callback": {
"get": {
"description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.",

135
tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py
View file

@ -97,6 +97,88 @@ def unauthenticated_client():
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("byok_first", [True, False])
def test_byok_challenge_discovers_api_key_flow(monkeypatch, byok_first):
from litellm.proxy._experimental.mcp_server import byok_oauth_endpoints, discoverable_endpoints
from litellm.proxy._experimental.mcp_server.oauth_utils import get_byok_www_authenticate
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
app = FastAPI()
routers = (byok_oauth_endpoints.router, discoverable_endpoints.router)
for item in routers if byok_first else reversed(routers):
app.include_router(item)
with TestClient(app) as session:
challenge = get_byok_www_authenticate()
assert challenge == 'Bearer resource_metadata="/v1/mcp/oauth/protected-resource"'
response = session.get(challenge.split('"')[1])
assert response.status_code == 200
assert response.json() == {
"resource": "http://testserver",
"authorization_servers": ["http://testserver/v1/mcp/oauth"],
}
authorization = session.get("/.well-known/oauth-authorization-server/v1/mcp/oauth")
assert authorization.status_code == 200
metadata = authorization.json()
assert metadata["issuer"] == response.json()["authorization_servers"][0]
assert metadata["authorization_endpoint"] == "http://testserver/v1/mcp/oauth/authorize"
assert metadata["token_endpoint"] == "http://testserver/v1/mcp/oauth/token"
assert metadata["code_challenge_methods_supported"] == ["S256"]
@pytest.mark.parametrize(
("base_url", "root_path", "expected"),
[
("", "", "/v1/mcp/oauth/protected-resource"),
("", "/proxy", "/proxy/v1/mcp/oauth/protected-resource"),
("https://gateway.example.com/proxy", "/proxy", "https://gateway.example.com/proxy/v1/mcp/oauth/protected-resource"),
],
)
def test_byok_challenge_preserves_external_base(monkeypatch, base_url, root_path, expected):
from litellm.proxy._experimental.mcp_server.oauth_utils import get_byok_www_authenticate
monkeypatch.setenv("PROXY_BASE_URL", base_url)
monkeypatch.setenv("SERVER_ROOT_PATH", root_path)
assert get_byok_www_authenticate() == f'Bearer resource_metadata="{expected}"'
def test_byok_discovery_preserves_per_request_prefixes(monkeypatch):
from fastapi import FastAPI
from litellm.proxy._experimental.mcp_server.server import _check_byok_credential
from litellm.proxy.middleware.per_request_root_path_middleware import PerRequestRootPathMiddleware
from litellm.types.mcp_server.mcp_server_manager import MCPServer
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
monkeypatch.setenv("SERVER_ROOT_PATHS", "/tenant-a,/tenant-b")
app = FastAPI()
app.include_router(router)
app.add_middleware(PerRequestRootPathMiddleware, root_paths=("/tenant-a", "/tenant-b"))
server = MCPServer(server_id="byok-prefix", name="byok-prefix", transport=MCPTransport.http, is_byok=True)
@app.get("/challenge")
async def challenge():
await _check_byok_credential(server, None)
with TestClient(app) as client:
for prefix in ("/tenant-a", "/tenant-b", ""):
challenge_response = client.get(f"{prefix}/challenge")
assert challenge_response.status_code == 401
metadata_path = f"{prefix}/v1/mcp/oauth/protected-resource"
assert challenge_response.headers["www-authenticate"] == f'Bearer resource_metadata="{metadata_path}"'
prm = client.get(metadata_path)
assert prm.status_code == 200
issuer = f"http://testserver{prefix}/v1/mcp/oauth"
assert prm.json()["authorization_servers"] == [issuer]
asm = client.get(f"/.well-known/oauth-authorization-server{prefix}/v1/mcp/oauth")
assert asm.status_code == 200
assert asm.json()["issuer"] == issuer
assert asm.json()["authorization_endpoint"] == f"{issuer}/authorize"
assert asm.json()["token_endpoint"] == f"{issuer}/token"
assert client.get("/.well-known/oauth-authorization-server/unknown/v1/mcp/oauth").status_code == 404
def test_oauth_authorization_server_metadata(client):
resp = client.get("/.well-known/oauth-authorization-server")
assert resp.status_code == 200
@ -107,15 +189,6 @@ def test_oauth_authorization_server_metadata(client):
assert "S256" in data["code_challenge_methods_supported"]
def test_oauth_protected_resource_metadata(client):
resp = client.get("/.well-known/oauth-protected-resource")
assert resp.status_code == 200
data = resp.json()
assert "resource" in data
assert "authorization_servers" in data
assert len(data["authorization_servers"]) == 1
# ---------------------------------------------------------------------------
# Authorization GET endpoint
# ---------------------------------------------------------------------------
@ -501,7 +574,7 @@ async def test_check_byok_credential_no_user_id():
@pytest.mark.asyncio
async def test_check_byok_credential_missing_credential():
async def test_check_byok_credential_missing_credential(monkeypatch):
"""BYOK server with a known user but no stored credential → 401."""
from litellm.proxy._experimental.mcp_server.server import _check_byok_credential
from litellm.proxy._types import UserAPIKeyAuth
@ -515,6 +588,11 @@ async def test_check_byok_credential_missing_credential():
)
user_auth = UserAPIKeyAuth(user_id="user-99", api_key="sk-test")
from litellm.proxy._experimental.mcp_server import server as server_module
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
monkeypatch.setattr(server_module, "_byok_cred_cache", {})
mock_prisma = MagicMock()
with (
@ -526,6 +604,10 @@ async def test_check_byok_credential_missing_credential():
):
with pytest.raises(HTTPException) as exc_info:
await _check_byok_credential(server, user_auth)
with pytest.raises(HTTPException) as cached_exc:
await _check_byok_credential(server, user_auth)
assert cached_exc.value.status_code == 401
assert cached_exc.value.headers == exc_info.value.headers
assert exc_info.value.status_code == 401
detail: Any = exc_info.value.detail
@ -533,7 +615,38 @@ async def test_check_byok_credential_missing_credential():
assert detail["server_id"] == "byok-2"
headers = exc_info.value.headers or {}
assert "WWW-Authenticate" in headers # type: ignore[operator]
assert "oauth-protected-resource" in headers["WWW-Authenticate"] # type: ignore[index]
assert headers["WWW-Authenticate"] == 'Bearer resource_metadata="/v1/mcp/oauth/protected-resource"'
@pytest.mark.asyncio
async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monkeypatch):
from datetime import datetime, timezone
from litellm.proxy._experimental.mcp_server import server as mcp_module
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy")
monkeypatch.setattr(mcp_module, "_byok_cred_cache", {})
server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True)
prisma = MagicMock()
prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
with pytest.raises(HTTPException) as exc_info:
await mcp_module.execute_mcp_tool(
name="list_regions",
arguments={},
allowed_mcp_servers=[server],
requested_server_id=server.server_id,
start_time=datetime.now(timezone.utc),
user_api_key_auth=UserAPIKeyAuth(user_id="byok-discovery-user"),
)
assert exc_info.value.status_code == 401
assert exc_info.value.detail["server_id"] == server.server_id
assert exc_info.value.headers == {
"WWW-Authenticate": 'Bearer resource_metadata="https://gateway.example.com/proxy/v1/mcp/oauth/protected-resource"'
}
@pytest.mark.asyncio

196
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -3433,51 +3433,84 @@ async def test_oauth_protected_resource_gateway_managed_oauth2_advertises_gatewa
global_mcp_server_manager.registry.clear()
@pytest.mark.parametrize("server_count", [0, 1, 2])
@pytest.mark.parametrize("byok_first", [True, False])
@pytest.mark.parametrize(
("base_url", "origin"),
[
("https://gateway.example.com", "https://gateway.example.com"),
("https://gateway.example.com/proxy", "https://gateway.example.com"),
("http://[::1]:4000/proxy", "http://[::1]:4000"),
],
)
def test_root_protected_resource_discovers_gateway(monkeypatch, server_count, byok_first, base_url, origin):
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy._experimental.mcp_server import byok_oauth_endpoints, discoverable_endpoints
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
monkeypatch.setenv("PROXY_BASE_URL", base_url)
monkeypatch.setattr(
global_mcp_server_manager,
"registry",
{
f"oauth_{index}": MCPServer(
server_id=f"oauth_{index}",
name=f"oauth_{index}",
server_name=f"oauth_{index}",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
authorization_url="https://idp.example.com/authorize",
token_url="https://idp.example.com/token",
)
for index in range(server_count)
},
)
app = FastAPI()
routers = (byok_oauth_endpoints.router, discoverable_endpoints.router)
for router in routers if byok_first else reversed(routers):
app.include_router(router)
with TestClient(app) as client:
response = client.get("/.well-known/oauth-protected-resource", params={"mcp_server_name": "oauth_0"})
assert response.status_code == 200
assert response.json() == {
"resource": origin,
"authorization_servers": [f"{base_url}/mcp"],
"scopes_supported": [],
}
authorization = client.get("/.well-known/oauth-authorization-server/mcp")
assert authorization.status_code == 200
metadata = authorization.json()
assert metadata["issuer"] == response.json()["authorization_servers"][0]
assert metadata["authorization_endpoint"] == f"{base_url}/authorize/mcp-session"
assert metadata["token_endpoint"] == f"{base_url}/token"
assert metadata["registration_endpoint"] == f"{base_url}/register"
aggregate = client.get("/.well-known/oauth-protected-resource/mcp")
assert aggregate.status_code == 200
assert aggregate.json()["resource"] == f"{base_url}/mcp"
@pytest.mark.asyncio
async def test_oauth_protected_resource_root_resolved_single_server_keeps_relay_as():
"""The unnamed (bare-root) legacy shape resolves the single configured oauth2 server and
must keep advertising the per-server relay authorization server: only an EXPLICITLY
named request opts into the gateway-as-AS flow (LIT-4864), so pre-existing single-server
deployments discovering through the root document are byte-identical."""
try:
from fastapi import Request
async def test_unnamed_protected_resource_builder_uses_gateway_origin(monkeypatch):
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_build_oauth_protected_resource_response,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
only_server = MCPServer(
server_id="solo_mcp",
name="solo_mcp",
server_name="solo_mcp",
alias="solo_mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
authorization_url="https://idp.example.com/authorize",
token_url="https://idp.example.com/oauth/token",
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_build_oauth_protected_resource_response,
)
mock_request = MagicMock(spec=Request)
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
global_mcp_server_manager.registry.clear()
try:
global_mcp_server_manager.registry[only_server.server_id] = only_server
response = await _build_oauth_protected_resource_response(
request=mock_request, mcp_server_name=None, use_standard_pattern=False
)
assert response["authorization_servers"] == ["https://litellm.example.com/solo_mcp"]
finally:
global_mcp_server_manager.registry.clear()
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
request = Request(
{"type": "http", "scheme": "https", "server": ("gateway.example.com", 443), "path": "/", "headers": []}
)
response = await _build_oauth_protected_resource_response(request, None, False)
assert response == {
"resource": "https://gateway.example.com",
"authorization_servers": ("https://gateway.example.com/mcp",),
"scopes_supported": (),
}
@pytest.mark.asyncio
@ -4137,8 +4170,9 @@ async def test_discovery_root_does_not_expose_private_server_for_external_client
assert "/test_oauth/" not in authorization_response["authorization_endpoint"]
assert "/test_oauth/" not in authorization_response["token_endpoint"]
assert authorization_response["scopes_supported"] == []
assert resource_response["authorization_servers"] == ["https://llm.example.com"]
assert resource_response["scopes_supported"] == []
assert tuple(resource_response["authorization_servers"]) == ("https://llm.example.com/mcp",)
assert resource_response["resource"] == "https://llm.example.com"
assert not resource_response["scopes_supported"]
finally:
global_mcp_server_manager.registry.clear()
@ -9014,7 +9048,7 @@ def test_aggregate_wellknown_routes_serve_gateway_metadata():
assert asm.status_code == 200
assert asm.json()["issuer"] == "http://testserver/mcp"
assert asm.json()["authorization_endpoint"] == "http://testserver/authorize"
assert asm.json()["authorization_endpoint"] == "http://testserver/authorize/mcp-session"
assert "none" in asm.json()["token_endpoint_auth_methods_supported"]
@ -9077,11 +9111,7 @@ def test_well_known_root_suffix_reflects_server_root_path():
@pytest.mark.asyncio
async def test_bare_origin_discovery_resolves_single_server_not_aggregate():
"""The always-on aggregate front door must not change bare-origin discovery: with one
oauth2 server configured, the no-suffix /.well-known/oauth-{authorization-server,
protected-resource} still resolves THAT server, so an existing single-server deployment's
discovery is unchanged. The aggregate document lives only at the /mcp-suffixed routes."""
async def test_root_resource_uses_gateway_without_changing_authorization_relay():
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
@ -9105,10 +9135,10 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate():
resource_response = await _build_oauth_protected_resource_response(
request=mock_request, mcp_server_name=None, use_standard_pattern=True
)
# per-server, not aggregate: the single server's name is in the endpoints
assert "/test_oauth/authorize" in authorization_response["authorization_endpoint"]
assert authorization_response["issuer"] == "https://llm.example.com"
assert resource_response["authorization_servers"] == ["https://llm.example.com/test_oauth"]
assert tuple(resource_response["authorization_servers"]) == ("https://llm.example.com/mcp",)
assert resource_response["resource"] == "https://llm.example.com"
finally:
global_mcp_server_manager.registry.clear()
@ -10640,7 +10670,7 @@ class TestPerRequestRootPathDiscovery:
assert asm.status_code == 200
assert asm.json()["issuer"] == "http://testserver/tenant-a/mcp"
assert asm.json()["authorization_endpoint"] == "http://testserver/tenant-a/authorize"
assert asm.json()["authorization_endpoint"] == "http://testserver/tenant-a/authorize/mcp-session"
# The prefixed authorize URL routes to the real handler (not 404):
# under per-request root_path the whole app is reachable per-prefix,
@ -10799,6 +10829,68 @@ def _consent_flow_handle(page: str) -> str:
return match.group(1)
@pytest.mark.parametrize("redirect_uri", ["http://127.0.0.1:51234/callback", "https://client.example.com/callback"])
@pytest.mark.parametrize("signed_in", [True, False])
def test_root_discovery_origin_authorizes_mcp_session(monkeypatch, redirect_uri, signed_in):
from urllib.parse import parse_qs, urlparse
client, session_cookie, minted = _native_client_app(monkeypatch)
root = client.get("/.well-known/oauth-protected-resource")
assert root.status_code == 200
assert root.json()["resource"] == "http://testserver"
authorization = client.get("/.well-known/oauth-authorization-server/mcp")
assert authorization.status_code == 200
metadata = authorization.json()
registered = client.post(metadata["registration_endpoint"], json={"redirect_uris": [redirect_uri]})
assert registered.status_code == 201
if signed_in:
client.cookies.set("token", session_cookie)
response = client.get(
metadata["authorization_endpoint"],
params={
"response_type": "code",
"client_id": registered.json()["client_id"],
"redirect_uri": redirect_uri,
"state": "mcp-state",
"code_challenge": _s256("v" * 43),
"code_challenge_method": "S256",
"resource": root.json()["resource"],
},
follow_redirects=False,
)
assert response.status_code == 303
target = urlparse(response.headers["location"])
assert target.path == ("/ui/connect" if signed_in else "/sso/key/generate")
if signed_in:
flow = parse_qs(target.query)["connect_flow"][0]
described = client.get("/authorize/flow", params={"flow": flow})
assert described.status_code == 200
assert described.json()["state"] == "unscoped"
assert minted == []
@pytest.mark.parametrize("valid_client", [True, False])
def test_mcp_session_authorize_rejects_invalid_registration_or_pkce(monkeypatch, valid_client):
client, session_cookie, minted = _native_client_app(monkeypatch)
redirect_uri = "https://client.example.com/callback"
registered = client.post("/register", json={"redirect_uris": [redirect_uri]})
assert registered.status_code == 201
client.cookies.set("token", session_cookie)
response = client.get(
"/authorize/mcp-session",
params={
"client_id": registered.json()["client_id"] if valid_client else "unknown-client",
"redirect_uri": redirect_uri,
"response_type": "code",
},
follow_redirects=False,
)
assert response.status_code == 400
assert response.json()["error"] == ("invalid_request" if valid_client else "invalid_client")
assert "location" not in response.headers
assert minted == []
def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(monkeypatch):
"""The whole ``lite login --pkce`` server side over the real router: a Go CLI reads the versioned
discovery document, registers a loopback public client, the signed-in user consents to a team,

6
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py
View file

@ -1233,13 +1233,14 @@ class TestResolveByokMcpAuthHeader:
assert result == "stored-cred"
@pytest.mark.asyncio
async def test_byok_server_raises_401_when_no_credential_stored(self):
async def test_byok_server_raises_401_when_no_credential_stored(self, monkeypatch):
from fastapi import HTTPException
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
_resolve_byok_mcp_auth_header,
)
monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy")
server = self._server(is_byok=True)
user_auth = UserAPIKeyAuth(user_id="user-1", api_key="sk-dashboard")
@ -1252,6 +1253,9 @@ class TestResolveByokMcpAuthHeader:
assert exc_info.value.status_code == 401
assert exc_info.value.detail["error"] == "byok_auth_required"
assert exc_info.value.headers == {
"WWW-Authenticate": 'Bearer resource_metadata="https://gateway.example.com/proxy/v1/mcp/oauth/protected-resource"'
}
@pytest.mark.asyncio
async def test_byok_server_checks_credential_and_keeps_caller_header_when_supplied(self):

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

@ -207,17 +207,8 @@ export interface paths {
path?: never;
cookie?: never;
};
/**
* Oauth Protected Resource Mcp
* @description OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.
*
* Legacy pattern: /{server_name}/mcp
* Discovery path: /.well-known/oauth-protected-resource/{server_name}/mcp
*
* This endpoint is kept for backward compatibility. New integrations should
* use the standard MCP pattern (/mcp/{server_name}) instead.
*/
get: operations["oauth_protected_resource_mcp__well_known_oauth_protected_resource_get"];
/** Oauth Protected Resource Root */
get: operations["oauth_protected_resource_root__well_known_oauth_protected_resource_get"];
put?: never;
post?: never;
delete?: never;
@ -1186,6 +1177,23 @@ export interface paths {
patch?: never;
trace?: never;
};
"/authorize/mcp-session": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
/** Authorize Mcp Session */
get: operations["authorize_mcp_session_authorize_mcp_session_get"];
put?: never;
post?: never;
delete?: never;
options?: never;
head?: never;
patch?: never;
trace?: never;
};
"/auto_router/benchmarks": {
parameters: {
query?: never;
@ -40285,11 +40293,9 @@ export interface operations {
};
};
};
oauth_protected_resource_mcp__well_known_oauth_protected_resource_get: {
oauth_protected_resource_root__well_known_oauth_protected_resource_get: {
parameters: {
query?: {
mcp_server_name?: string | null;
};
query?: never;
header?: never;
path?: never;
cookie?: never;
@ -40302,16 +40308,9 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": unknown;
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
"application/json": {
[key: string]: string | string[];
};
};
};
};
@ -41666,6 +41665,43 @@ export interface operations {
};
};
};
authorize_mcp_session_authorize_mcp_session_get: {
parameters: {
query: {
redirect_uri: string;
client_id: string;
state?: string;
code_challenge?: string | null;
code_challenge_method?: string | null;
response_type?: string | null;
resource?: string | null;
};
header?: never;
path?: never;
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": unknown;
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
get_auto_router_benchmarks_auto_router_benchmarks_get: {
parameters: {
query?: {