scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

Merge pull request #26274 from stuxf/fix/mcp-byok-oauth-auth

Browse source 
fix(mcp): harden OAuth authorize/token endpoints (BYOK + discoverable)
This commit is contained in:
yuneng-jiang 2026-04-24 13:05:44 -07:00 • committed by GitHub
commit 51d4c6c2f2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1085 additions and 116 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -19,10 +19,10 @@ import html as _html_module
import time
import uuid
from typing import Dict, Optional, cast
from urllib.parse import urlencode, urlparse
from urllib.parse import urlencode
import jwt
from fastapi import APIRouter, Form, HTTPException, Request
from fastapi import APIRouter, Depends, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
from litellm._logging import verbose_proxy_logger
@ -30,6 +30,11 @@ 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 (
TOKEN_NO_CACHE_HEADERS,
validate_loopback_redirect_uri,
)
from litellm.proxy._types import UserAPIKeyAuth
# ---------------------------------------------------------------------------
# In-memory store for pending authorization codes.
@ -69,6 +74,65 @@ def _purge_expired_codes() -> None:
del _byok_auth_codes[k]
def _oauth_token_error(code: str, status: int = 400) -> JSONResponse:
"""RFC 6749 §5.2 token-endpoint error body: ``{"error": "<code>"}``.
FastAPI's default ``HTTPException`` renders ``{"detail": ...}`` which
spec-compliant OAuth clients parsing the ``error`` field won't recognize.
"""
return JSONResponse(
status_code=status, content={"error": code}, headers=TOKEN_NO_CACHE_HEADERS
)
def _user_id_from_session_cookie(request: Request) -> Optional[str]:
"""Return user_id from the UI ``token`` cookie (HS256-signed with
``master_key``), or None if missing/invalid.
The /token endpoint in this file ALSO issues master-key-signed JWTs
(type="byok_session") for MCP-client-side use. They must not be
accepted here as UI sessions — otherwise a leaked byok_session token
could be replayed as a cookie to re-authorize BYOK writes. Distinguish
by requiring a ``login_method`` claim (UI tokens set ``"sso"`` or
``"username_password"``; byok_session tokens never set it) and
rejecting any token whose ``type`` identifies it as non-UI.
"""
# Inline import avoids a circular dep (proxy_server -> mcp_server router).
from litellm.proxy.proxy_server import master_key
if not master_key:
return None
token = request.cookies.get("token")
if not token:
return None
try:
payload = jwt.decode(
token,
master_key,
algorithms=["HS256"],
# Require an expiry claim so a leaked UI session cookie has a
# bounded lifetime. PyJWT verifies exp by default when present;
# require=["exp"] additionally rejects tokens that omit it.
options={"require": ["exp"]},
)
except jwt.InvalidTokenError:
return None
if payload.get("type") == "byok_session":
return None
if payload.get("login_method") not in ("sso", "username_password"):
return None
user_id = payload.get("user_id")
return user_id if isinstance(user_id, str) and user_id else None
async def _byok_session_auth(request: Request) -> UserAPIKeyAuth:
"""Require the UI session cookie. Programmatic BYOK management uses
``POST /v1/mcp/server/{id}/user-credential`` instead."""
user_id = _user_id_from_session_cookie(request)
if not user_id:
raise HTTPException(status_code=401, detail="login_required")
return UserAPIKeyAuth(api_key="byok_session_cookie", user_id=user_id)
def _build_authorize_html(
server_name: str,
server_initial: str,
@ -582,11 +646,18 @@ async def byok_authorize_get(
The MCP client navigates the user here; the user types their API key and
clicks "Connect & Authorize", which POSTs back to this same path.
This GET is intentionally unauthenticated: it only renders HTML with no
state change. The POST handler enforces ``user_api_key_auth`` and pins
the stored credential to the authenticated session.
"""
if response_type != "code":
raise HTTPException(status_code=400, detail="response_type must be 'code'")
if not redirect_uri:
raise HTTPException(status_code=400, detail="redirect_uri is required")
# Validate here too so the user sees the rejection before typing their
# API key into the HTML form (the POST handler also validates).
validate_loopback_redirect_uri(redirect_uri)
if not code_challenge:
raise HTTPException(status_code=400, detail="code_challenge is required")
@ -636,6 +707,7 @@ async def byok_authorize_post(
state: str = Form(default=""),
server_id: str = Form(default=""),
api_key: str = Form(...),
user_api_key_dict: UserAPIKeyAuth = Depends(_byok_session_auth),
) -> RedirectResponse:
"""
Process the BYOK API-key form submission.
@ -645,10 +717,7 @@ async def byok_authorize_post(
"""
_purge_expired_codes()
# Validate redirect_uri scheme to prevent open redirect
parsed_uri = urlparse(redirect_uri)
if parsed_uri.scheme not in ("http", "https"):
raise HTTPException(status_code=400, detail="Invalid redirect_uri scheme")
validate_loopback_redirect_uri(redirect_uri)
# Reject new codes if the store is at capacity (prevents memory exhaustion
# from a burst of abandoned OAuth flows).
@ -662,13 +731,25 @@ async def byok_authorize_post(
status_code=400, detail="Only S256 code_challenge_method is supported"
)
# Identity comes from the authenticated session, not the OAuth client_id
# form field (RFC 6749 §2.2: client_id identifies the client application,
# not the user). We do bind the code to the submitted client_id so the
# /token call must present the same value (RFC 6749 §4.1.3).
user_id = user_api_key_dict.user_id
if not user_id:
raise HTTPException(status_code=401, detail="login_required")
auth_code = str(uuid.uuid4())
_byok_auth_codes[auth_code] = {
"api_key": api_key,
"server_id": server_id,
"code_challenge": code_challenge,
"redirect_uri": redirect_uri,
"user_id": client_id, # external client passes LiteLLM user-id as client_id
# RFC 6749 §4.1.3 defense-in-depth: if the authorization request
# declared a client_id, the token request must submit the same
# value. Stored even though we don't pre-register clients.
"client_id": client_id,
"user_id": user_id,
"expires_at": time.time() + _AUTH_CODE_TTL_SECONDS,
}
@ -704,34 +785,60 @@ async def byok_token(
_purge_expired_codes()
if grant_type != "authorization_code":
raise HTTPException(status_code=400, detail="unsupported_grant_type")
return _oauth_token_error("unsupported_grant_type")
record = _byok_auth_codes.get(code)
if record is None:
raise HTTPException(status_code=400, detail="invalid_grant")
return _oauth_token_error("invalid_grant")
if time.time() > record["expires_at"]:
del _byok_auth_codes[code]
raise HTTPException(status_code=400, detail="invalid_grant")
return _oauth_token_error("invalid_grant")
# PKCE verification
if not _verify_pkce(code_verifier, record["code_challenge"]):
raise HTTPException(status_code=400, detail="invalid_grant")
return _oauth_token_error("invalid_grant")
# Consume the code (one-time use)
del _byok_auth_codes[code]
# RFC 6749 §4.1.3: if redirect_uri was sent with the authorization
# request, the token request MUST include the identical value.
# OAuth 2.1 draft-15 §4.1.3 drops this requirement — strict OAuth 2.1
# clients will omit it. Enforce equality ONLY when the client
# actually submitted a value, so we stay RFC 6749-backward-compatible
# without breaking OAuth 2.1 clients. PKCE + client_id binding
# (checked below) cover the security role redirect_uri played.
if (
record.get("redirect_uri")
and redirect_uri
and redirect_uri != record["redirect_uri"]
):
return _oauth_token_error("invalid_grant")
# RFC 6749 §4.1.3: if the client was identified at /authorize, the
# /token request MUST authenticate as the same client. We don't
# pre-register clients, so an empty stored client_id skips the check.
if record.get("client_id") and client_id != record["client_id"]:
return _oauth_token_error("invalid_grant")
server_id: str = record["server_id"]
api_key_value: str = record["api_key"]
# Prefer the user_id that was stored when the code was issued; fall back to
# whatever client_id the token request supplies (they should match).
user_id: str = record.get("user_id") or client_id
# user_id is stamped by the authenticated /authorize POST. No client_id
# fallback — that fallback was the credential-hijack primitive. The
# token-endpoint client_id is informational per RFC 6749 and is not
# cross-checked against user_id (which identifies the resource owner,
# not the client application).
user_id: str = record.get("user_id") or ""
if not user_id:
raise HTTPException(
status_code=400,
detail="Cannot determine user_id; pass LiteLLM user id as client_id",
)
return _oauth_token_error("invalid_grant")
# Verify preconditions that would fail token issuance BEFORE consuming
# the code or writing to the DB — otherwise a misconfigured proxy
# (missing master_key) silently persists the user's credential without
# ever returning an access token, and the user has no way to recover.
if master_key is None:
return _oauth_token_error("server_error", status=500)
# Consume the code (one-time use)
del _byok_auth_codes[code]
# Persist the BYOK credential
if prisma_client is not None:
@ -756,17 +863,12 @@ async def byok_token(
server_id,
exc,
)
raise HTTPException(status_code=500, detail="Failed to store credential")
return _oauth_token_error("server_error", status=500)
else:
verbose_proxy_logger.warning(
"byok_token: prisma_client is None — credential not persisted"
)
if master_key is None:
raise HTTPException(
status_code=500, detail="Master key not configured; cannot issue token"
)
now = int(time.time())
payload = {
"user_id": user_id,
@ -785,5 +887,6 @@ async def byok_token(
"access_token": access_token,
"token_type": "bearer",
"expires_in": 3600,
}
},
headers=TOKEN_NO_CACHE_HEADERS,
)

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

@ -10,6 +10,10 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._experimental.mcp_server.oauth_utils import (
TOKEN_NO_CACHE_HEADERS,
validate_loopback_redirect_uri,
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
@ -322,15 +326,13 @@ async def authorize_with_server(
status_code=400, detail="MCP server authorization url is not set"
)
# Loopback-only redirect_uri. The URI is encrypted into the OAuth
# state and decoded on /callback to redirect the user back; a non-
# loopback URI would be an open-redirect + code-theft primitive
# (VERIA-57 root cause B). MCP clients are native apps — loopback is
# the spec-compliant callback pattern.
validate_loopback_redirect_uri(redirect_uri)
parsed = urlparse(redirect_uri)
if parsed.scheme not in ("http", "https"):
raise HTTPException(
status_code=400,
detail={
"error": "invalid_redirect_uri",
"message": "redirect_uri must use http or https scheme",
},
)
base_url = urlunparse(parsed._replace(query=""))
request_base_url = get_request_base_url(request)
encoded_state = encode_state_with_base_url(
@ -480,7 +482,8 @@ async def exchange_token_with_server(
if "scope" in token_response and token_response["scope"]:
result["scope"] = token_response["scope"]
return JSONResponse(result)
# RFC 6749 §5.1: token responses must not be cached.
return JSONResponse(result, headers=TOKEN_NO_CACHE_HEADERS)
async def register_client_with_server(
@ -647,20 +650,26 @@ async def token_endpoint(
@router.get("/callback")
async def callback(code: str, state: str):
try:
# Decode the state hash to get base_url, original state, and PKCE params
state_data = decode_state_hash(state)
base_url = state_data["base_url"]
original_state = state_data["original_state"]
# Forward code and original state back to client
params = {"code": code, "state": original_state}
# Re-validate loopback at the sink. /authorize rejects non-loopback
# redirect_uri before encoding into state, but encrypted states
# minted before that check was added have no expiry and remain
# valid indefinitely. Validating here blocks the open-redirect +
# code-theft primitive even for pre-fix states.
validate_loopback_redirect_uri(base_url)
# Forward to client's callback endpoint
params = {"code": code, "state": original_state}
complete_returned_url = f"{base_url}?{urlencode(params)}"
return RedirectResponse(url=complete_returned_url, status_code=302)
except HTTPException:
# Re-raise so a non-loopback base_url surfaces as 400 instead of
# a generic "authentication incomplete" redirect.
raise
except Exception:
# fallback if state hash not found
return HTMLResponse(
"<html><body>Authentication incomplete. You can close this window.</body></html>"
)

48
litellm/proxy/_experimental/mcp_server/oauth_utils.py Normal file
View file

@ -0,0 +1,48 @@
"""Shared helpers for the MCP OAuth authorization endpoints
(BYOK + discoverable / pass-through OAuth proxy)."""
from ipaddress import ip_address
from urllib.parse import urlparse
from fastapi import HTTPException
# RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: token-endpoint responses
# must not be cached — both success and error bodies may reveal secrets.
TOKEN_NO_CACHE_HEADERS = {"Cache-Control": "no-store", "Pragma": "no-cache"}
def validate_loopback_redirect_uri(redirect_uri: str) -> None:
"""Require a loopback ``redirect_uri`` (OAuth 2.1 §4.1.2.1 + RFC 8252
§7.3 native-app pattern). MCP clients are native apps that listen on
a localhost port; rejecting non-loopback URIs prevents a malicious
client from pointing the callback at its own server to capture the
authorization code — the credential-theft primitive behind VERIA-57
and pNr1PHa9.
Accepts the literal ``localhost`` plus any IP in the loopback ranges
(IPv4 ``127.0.0.0/8`` and IPv6 ``::1``). A string match on
``"127.0.0.1"`` alone would miss ``127.0.0.2`` and the full-form
IPv6 loopback ``0:0:0:0:0:0:0:1``.
"""
try:
parsed = urlparse(redirect_uri)
except ValueError:
raise HTTPException(status_code=400, detail="invalid_request")
if parsed.scheme not in ("http", "https"):
raise HTTPException(status_code=400, detail="invalid_request")
# Fragments are not allowed in OAuth redirect URIs (RFC 6749 §3.1.2)
# — rejecting them prevents a ``http://127.0.0.1/cb#frag?code=...``
# from silently eating the authorization code.
if parsed.fragment:
raise HTTPException(status_code=400, detail="invalid_request")
host = (parsed.hostname or "").lower()
if host == "localhost":
return
try:
if ip_address(host).is_loopback:
return
except ValueError:
# Unparseable host (malformed IPv6, etc.) — treat as invalid,
# don't let it bubble up as a 500.
pass
raise HTTPException(status_code=400, detail="invalid_request")

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

@ -1952,7 +1952,18 @@ if MCP_AVAILABLE:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return
# Fail closed on DB unavailability: returning here previously
# bypassed the ownership check and let any proxy-authenticated
# caller invoke BYOK tools during outage windows.
raise HTTPException(
status_code=503,
detail={
"error": "byok_auth_unavailable",
"server_id": mcp_server.server_id,
"server_name": mcp_server.server_name or mcp_server.name,
"message": "BYOK credential check requires a database connection.",
},
)
credential = await get_user_credential(
prisma_client=prisma_client,

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

@ -11,6 +11,7 @@ Covers:
import base64
import hashlib
import json
import time
import uuid
from typing import Any, Optional
@ -61,13 +62,34 @@ def test_verify_pkce_tampered_challenge():
from fastapi import FastAPI
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
_byok_session_auth,
)
from litellm.proxy._types import UserAPIKeyAuth
_test_app = FastAPI()
_test_app.include_router(router)
@pytest.fixture
def client():
return TestClient(_test_app, raise_server_exceptions=False)
"""Test client with a fixed authenticated user (bypasses the session
cookie check by overriding the dep)."""
_test_app.dependency_overrides[_byok_session_auth] = lambda: UserAPIKeyAuth(
api_key="hashed", user_id="user-123"
)
try:
yield TestClient(_test_app, raise_server_exceptions=False)
finally:
_test_app.dependency_overrides.pop(_byok_session_auth, None)
@pytest.fixture
def unauthenticated_client():
"""Test client with no dependency override — the real ``_byok_session_auth``
runs, which checks the ``token`` cookie and falls back to
``user_api_key_auth``. With neither set, both paths fail → 401."""
yield TestClient(_test_app, raise_server_exceptions=False)
# ---------------------------------------------------------------------------
@ -104,7 +126,7 @@ def test_authorize_get_returns_html(client):
"/v1/mcp/oauth/authorize",
params={
"client_id": "test-client",
"redirect_uri": "https://client.example.com/callback",
"redirect_uri": "http://127.0.0.1:3000/callback",
"response_type": "code",
"code_challenge": "abc123",
"code_challenge_method": "S256",
@ -138,7 +160,7 @@ def test_authorize_get_wrong_response_type(client):
resp = client.get(
"/v1/mcp/oauth/authorize",
params={
"redirect_uri": "https://example.com/cb",
"redirect_uri": "http://127.0.0.1:3000/cb",
"response_type": "token",
"code_challenge": "abc",
},
@ -147,6 +169,23 @@ def test_authorize_get_wrong_response_type(client):
assert resp.status_code == 400
def test_authorize_get_rejects_non_loopback_redirect_uri(client):
"""GET /authorize validates redirect_uri up front so the user sees
the rejection before typing an API key into the HTML form — matches
the POST handler's rule and avoids the ``user fills form → POST 400
with no form state`` UX."""
resp = client.get(
"/v1/mcp/oauth/authorize",
params={
"redirect_uri": "https://attacker.example.com/cb",
"response_type": "code",
"code_challenge": "abc",
},
follow_redirects=False,
)
assert resp.status_code == 400
# ---------------------------------------------------------------------------
# Authorization POST endpoint
# ---------------------------------------------------------------------------
@ -160,7 +199,7 @@ def test_authorize_post_creates_code_and_redirects(client):
"/v1/mcp/oauth/authorize",
data={
"client_id": "user-123",
"redirect_uri": "https://client.example.com/callback",
"redirect_uri": "http://127.0.0.1:3000/callback",
"code_challenge": challenge,
"code_challenge_method": "S256",
"state": "st_abc",
@ -286,10 +325,11 @@ async def test_token_endpoint_success():
)
assert result.status_code == 200
body = result.body
import json
data = json.loads(body)
# RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: token responses MUST NOT
# be cached, as the body contains an access token.
assert result.headers["cache-control"] == "no-store"
assert result.headers["pragma"] == "no-cache"
data = json.loads(result.body)
assert "access_token" in data
assert data["token_type"] == "bearer"
assert data["expires_in"] == 3600
@ -319,21 +359,20 @@ async def test_token_endpoint_invalid_code():
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import byok_token
mock_request = MagicMock()
with pytest.raises(HTTPException) as exc_info:
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "key"),
):
await byok_token(
request=mock_request,
grant_type="authorization_code",
code="nonexistent-code",
redirect_uri="",
code_verifier="anything",
client_id="u",
)
assert exc_info.value.status_code == 400
assert "invalid_grant" in str(exc_info.value.detail)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "key"),
):
result = await byok_token(
request=mock_request,
grant_type="authorization_code",
code="nonexistent-code",
redirect_uri="",
code_verifier="anything",
client_id="u",
)
assert result.status_code == 400
assert json.loads(result.body) == {"error": "invalid_grant"}
@pytest.mark.asyncio
@ -345,27 +384,27 @@ async def test_token_endpoint_expired_code():
server_id="s",
user_id="u",
challenge=challenge,
redirect_uri="https://cb",
redirect_uri="http://127.0.0.1:3000/cb",
ttl=-10, # already expired
)
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import byok_token
mock_request = MagicMock()
with pytest.raises(HTTPException) as exc_info:
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "key"),
):
await byok_token(
request=mock_request,
grant_type="authorization_code",
code=code,
redirect_uri="",
code_verifier=verifier,
client_id="u",
)
assert exc_info.value.status_code == 400
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "key"),
):
result = await byok_token(
request=mock_request,
grant_type="authorization_code",
code=code,
redirect_uri="http://127.0.0.1:3000/cb",
code_verifier=verifier,
client_id="u",
)
assert result.status_code == 400
assert json.loads(result.body) == {"error": "invalid_grant"}
@pytest.mark.asyncio
@ -377,27 +416,26 @@ async def test_token_endpoint_wrong_verifier():
server_id="s",
user_id="u",
challenge=challenge,
redirect_uri="https://cb",
redirect_uri="http://127.0.0.1:3000/cb",
)
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import byok_token
mock_request = MagicMock()
with pytest.raises(HTTPException) as exc_info:
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "key"),
):
await byok_token(
request=mock_request,
grant_type="authorization_code",
code=code,
redirect_uri="",
code_verifier="wrong_verifier_value_that_wont_match",
client_id="u",
)
assert exc_info.value.status_code == 400
assert "invalid_grant" in str(exc_info.value.detail)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "key"),
):
result = await byok_token(
request=mock_request,
grant_type="authorization_code",
code=code,
redirect_uri="http://127.0.0.1:3000/cb",
code_verifier="wrong_verifier_value_that_wont_match",
client_id="u",
)
assert result.status_code == 400
assert json.loads(result.body) == {"error": "invalid_grant"}
@pytest.mark.asyncio
@ -405,21 +443,20 @@ async def test_token_endpoint_unsupported_grant_type():
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import byok_token
mock_request = MagicMock()
with pytest.raises(HTTPException) as exc_info:
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "key"),
):
await byok_token(
request=mock_request,
grant_type="client_credentials",
code="any",
redirect_uri="",
code_verifier="v",
client_id="u",
)
assert exc_info.value.status_code == 400
assert "unsupported_grant_type" in str(exc_info.value.detail)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", "key"),
):
result = await byok_token(
request=mock_request,
grant_type="client_credentials",
code="any",
redirect_uri="",
code_verifier="v",
client_id="u",
)
assert result.status_code == 400
assert json.loads(result.body) == {"error": "unsupported_grant_type"}
# ---------------------------------------------------------------------------
@ -525,3 +562,576 @@ async def test_check_byok_credential_has_credential():
):
# Should not raise
await _check_byok_credential(server, user_auth)
@pytest.mark.asyncio
async def test_check_byok_credential_db_unavailable_fails_closed():
"""BYOK server with no prisma_client → 503, not silent pass.
Regression for GHSA-6762: previously returned silently, bypassing the
ownership check during DB outage windows.
"""
from litellm.proxy._experimental.mcp_server.server import _check_byok_credential
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="byok-4",
name="byok-server",
transport=MCPTransport.http,
is_byok=True,
)
user_auth = UserAPIKeyAuth(user_id="user-55", api_key="sk-test")
with patch("litellm.proxy.proxy_server.prisma_client", None):
with pytest.raises(HTTPException) as exc_info:
await _check_byok_credential(server, user_auth)
assert exc_info.value.status_code == 503
detail: Any = exc_info.value.detail
assert detail["error"] == "byok_auth_unavailable"
assert detail["server_id"] == "byok-4"
# ---------------------------------------------------------------------------
# Security regression tests for AO2kf_-9 / GHSA-jg3h:
# Unauthenticated /v1/mcp/oauth/authorize previously allowed an attacker to
# stamp `user_id = client_id` into the auth-code record, overwriting any
# victim's stored BYOK credential at /token.
# ---------------------------------------------------------------------------
def test_authorize_post_rejects_unauthenticated(unauthenticated_client):
resp = unauthenticated_client.post(
"/v1/mcp/oauth/authorize",
data={
"client_id": "victim-user-id",
"redirect_uri": "http://127.0.0.1:3000/cb",
"code_challenge": "abc",
"code_challenge_method": "S256",
"state": "s",
"server_id": "sid",
"api_key": "attacker-key",
},
follow_redirects=False,
)
assert resp.status_code == 401
def test_authorize_post_ignores_spec_compliant_client_id(client):
"""An MCP client following OAuth 2.1 semantics sends client_id as its
*application* identifier (e.g. "claude-desktop"). The stored user_id
must come from the authenticated session regardless — the form's
client_id is informational only."""
verifier = "verifier_value_long_enough_to_be_valid_43chars"
challenge = _make_challenge(verifier)
resp = client.post(
"/v1/mcp/oauth/authorize",
data={
# Fixture authenticates as "user-123"; the client identifies
# itself as "claude-desktop" per OAuth 2.1 — unrelated to user.
"client_id": "claude-desktop",
"redirect_uri": "http://127.0.0.1:3000/cb",
"code_challenge": challenge,
"code_challenge_method": "S256",
"state": "s",
"server_id": "sid",
"api_key": "upstream-key",
},
follow_redirects=False,
)
assert resp.status_code == 302
from urllib.parse import parse_qs, urlparse
qs = parse_qs(urlparse(resp.headers["location"]).query)
code = qs["code"][0]
# Identity = authenticated session, not the form's client_id.
assert _byok_auth_codes[code]["user_id"] == "user-123"
def test_authorize_post_binds_code_to_authenticated_user_id(client):
"""Ensure the stored auth-code record uses the authenticated user_id,
NOT the form's client_id, as the identity the token endpoint will trust."""
verifier = "verifier_value_long_enough_to_be_valid_43chars"
challenge = _make_challenge(verifier)
resp = client.post(
"/v1/mcp/oauth/authorize",
data={
# client_id omitted — must still bind to the authenticated user.
"redirect_uri": "http://127.0.0.1:3000/cb",
"code_challenge": challenge,
"code_challenge_method": "S256",
"state": "s",
"server_id": "sid",
"api_key": "legit-key",
},
follow_redirects=False,
)
assert resp.status_code == 302
from urllib.parse import parse_qs, urlparse
qs = parse_qs(urlparse(resp.headers["location"]).query)
code = qs["code"][0]
assert _byok_auth_codes[code]["user_id"] == "user-123"
@pytest.mark.asyncio
async def test_token_endpoint_rejects_missing_user_id_in_code_record():
"""Defense in depth: if a code record somehow lacks user_id (older
format / manual DB write), /token must reject rather than fall back to
the form's client_id."""
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
byok_token,
)
verifier = "verifier_for_test_missing_user_id_path"
challenge = _make_challenge(verifier)
code = str(uuid.uuid4())
_byok_auth_codes[code] = {
"api_key": "k",
"server_id": "sid",
"code_challenge": challenge,
"redirect_uri": "http://127.0.0.1:3000/cb",
"user_id": "", # missing / empty
"expires_at": time.time() + 60,
}
result = await byok_token(
request=MagicMock(),
grant_type="authorization_code",
code=code,
redirect_uri="http://127.0.0.1:3000/cb",
code_verifier=verifier,
client_id="attacker-chosen",
)
assert result.status_code == 400
assert json.loads(result.body) == {"error": "invalid_grant"}
def _authorize_post_with_cookie(client, cookie_jwt: str, api_key: str = "upstream-key"):
verifier = "verifier_cookie_auth_long_enough_to_be_valid"
return client.post(
"/v1/mcp/oauth/authorize",
data={
"redirect_uri": "http://127.0.0.1:3000/cb",
"code_challenge": _make_challenge(verifier),
"code_challenge_method": "S256",
"state": "s",
"server_id": "sid",
"api_key": api_key,
},
cookies={"token": cookie_jwt},
follow_redirects=False,
)
def test_authorize_post_accepts_ui_session_cookie(unauthenticated_client):
"""Browser flow: the native HTML form doesn't add Authorization. Instead
the user's UI session cookie ``token`` carries a master-key-signed JWT
whose ``user_id`` + ``login_method`` claims authenticate the POST."""
import jwt as _jwt
with patch("litellm.proxy.proxy_server.master_key", "test-master-key"):
cookie_jwt = _jwt.encode(
{
"user_id": "browser-user-42",
"login_method": "sso",
"exp": int(time.time()) + 3600,
},
"test-master-key",
algorithm="HS256",
)
resp = _authorize_post_with_cookie(unauthenticated_client, cookie_jwt)
assert resp.status_code == 302
from urllib.parse import parse_qs, urlparse
qs = parse_qs(urlparse(resp.headers["location"]).query)
code = qs["code"][0]
assert _byok_auth_codes[code]["user_id"] == "browser-user-42"
def test_authorize_post_rejects_cookie_signed_with_wrong_key(unauthenticated_client):
"""A cookie JWT signed with a different key than the proxy's master_key
must not grant access — otherwise an attacker who can forge a JWT
against any key could impersonate any user."""
import jwt as _jwt
with patch("litellm.proxy.proxy_server.master_key", "real-master-key"):
forged = _jwt.encode(
{"user_id": "victim-user", "login_method": "sso"},
"attacker-key",
algorithm="HS256",
)
resp = _authorize_post_with_cookie(unauthenticated_client, forged, api_key="k")
assert resp.status_code == 401
def test_authorize_post_rejects_replayed_byok_session_token(unauthenticated_client):
"""Regression: the /token endpoint itself issues master-key-signed JWTs
with ``type="byok_session"`` + ``user_id`` (for MCP-client use). Those
tokens must not be accepted here — otherwise an attacker with any
byok_session token could replay it as a Cookie and re-authorize BYOK
writes without a valid UI session."""
import jwt as _jwt
with patch("litellm.proxy.proxy_server.master_key", "real-master-key"):
byok_session = _jwt.encode(
{
"user_id": "any-user",
"server_id": "sid",
"type": "byok_session",
},
"real-master-key",
algorithm="HS256",
)
resp = _authorize_post_with_cookie(
unauthenticated_client, byok_session, api_key="k"
)
assert resp.status_code == 401
def test_authorize_post_rejects_non_loopback_redirect_uri(client):
"""OAuth 2.1 §4.1.2.1 + RFC 8252: native-app redirects must be loopback.
A public HTTPS callback from an MCP client would let that client capture
the issued code after a legitimate user enters their API key, so we
reject anything that isn't 127.0.0.1/localhost/::1."""
verifier = "verifier_non_loopback_redirect_uri_test_long"
resp = client.post(
"/v1/mcp/oauth/authorize",
data={
"redirect_uri": "https://attacker.example.com/cb",
"code_challenge": _make_challenge(verifier),
"code_challenge_method": "S256",
"state": "s",
"server_id": "sid",
"api_key": "k",
},
follow_redirects=False,
)
assert resp.status_code == 400
@pytest.mark.asyncio
async def test_token_endpoint_rejects_redirect_uri_mismatch():
"""RFC 6749 §4.1.3 / OAuth 2.1 §4.1.3: if redirect_uri was sent at
/authorize, the /token redirect_uri MUST match exactly."""
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
byok_token,
)
verifier = "verifier_for_redirect_mismatch_test_long"
challenge = _make_challenge(verifier)
code = str(uuid.uuid4())
_byok_auth_codes[code] = {
"api_key": "k",
"server_id": "sid",
"code_challenge": challenge,
"redirect_uri": "http://127.0.0.1:3000/cb",
"user_id": "u",
"expires_at": time.time() + 60,
}
result = await byok_token(
request=MagicMock(),
grant_type="authorization_code",
code=code,
redirect_uri="http://127.0.0.1:9999/cb", # different port
code_verifier=verifier,
client_id="",
)
assert result.status_code == 400
assert json.loads(result.body) == {"error": "invalid_grant"}
def test_authorize_post_rejects_cookie_missing_login_method(unauthenticated_client):
"""Defense in depth: a master-key-signed JWT with only ``user_id`` is not
a valid UI session (UI tokens always carry ``login_method``). Accepting
it would expand the cookie surface to include any master-key-signed
JWT in the system, which is exactly what the byok_session-replay
regression above protects against."""
import jwt as _jwt
with patch("litellm.proxy.proxy_server.master_key", "real-master-key"):
malformed = _jwt.encode(
{"user_id": "some-user"},
"real-master-key",
algorithm="HS256",
)
resp = _authorize_post_with_cookie(
unauthenticated_client, malformed, api_key="k"
)
assert resp.status_code == 401
@pytest.mark.asyncio
async def test_token_endpoint_accepts_spec_compliant_client_id():
"""Per OAuth 2.1, the /token client_id is the client application
identifier, not the user. It must not be cross-checked against the
record's user_id — that cross-check would break spec-compliant MCP
clients that pass e.g. client_id="claude-desktop"."""
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
byok_token,
)
verifier = "verifier_for_spec_compliant_long_enough_43chars"
challenge = _make_challenge(verifier)
code = str(uuid.uuid4())
_byok_auth_codes[code] = {
"api_key": "k",
"server_id": "sid",
"code_challenge": challenge,
"redirect_uri": "http://127.0.0.1:3000/cb",
"user_id": "real-authenticated-user",
"expires_at": time.time() + 60,
}
with (
patch(
"litellm.proxy._experimental.mcp_server.byok_oauth_endpoints.store_user_credential",
new=AsyncMock(),
),
patch("litellm.proxy.proxy_server.master_key", "test-master"),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
):
result = await byok_token(
request=MagicMock(),
grant_type="authorization_code",
code=code,
redirect_uri="http://127.0.0.1:3000/cb",
code_verifier=verifier,
client_id="claude-desktop", # OAuth app id, unrelated to user
)
# Token should be issued successfully.
assert result.status_code == 200
@pytest.mark.asyncio
async def test_token_endpoint_rejects_client_id_mismatch():
"""RFC 6749 §4.1.3: if the authorization request was bound to a
client_id, the token request must submit the same value. An attacker
who steals a code from another client (different native app) can't
redeem it."""
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
byok_token,
)
verifier = "verifier_for_client_id_mismatch_long_enough!"
challenge = _make_challenge(verifier)
code = str(uuid.uuid4())
_byok_auth_codes[code] = {
"api_key": "k",
"server_id": "sid",
"code_challenge": challenge,
"redirect_uri": "http://127.0.0.1:3000/cb",
"client_id": "legitimate-client",
"user_id": "u",
"expires_at": time.time() + 60,
}
result = await byok_token(
request=MagicMock(),
grant_type="authorization_code",
code=code,
redirect_uri="http://127.0.0.1:3000/cb",
code_verifier=verifier,
client_id="attacker-client",
)
assert result.status_code == 400
assert json.loads(result.body) == {"error": "invalid_grant"}
def test_authorize_post_accepts_ipv4_loopback_range(client):
"""RFC 8252 §7.3 / RFC 5735: ``127.0.0.0/8`` is loopback — a string
match on ``127.0.0.1`` would miss ``127.0.0.2`` and break clients that
pick a loopback alias."""
verifier = "verifier_for_127002_loopback_test_long_enough"
resp = client.post(
"/v1/mcp/oauth/authorize",
data={
"redirect_uri": "http://127.0.0.2:3000/cb",
"code_challenge": _make_challenge(verifier),
"code_challenge_method": "S256",
"state": "s",
"server_id": "sid",
"api_key": "k",
},
follow_redirects=False,
)
assert resp.status_code == 302
def test_authorize_post_accepts_ipv6_loopback_full_form(client):
"""RFC 4291: full IPv6 loopback ``0:0:0:0:0:0:0:1`` must be accepted
equivalently to ``::1``."""
verifier = "verifier_for_ipv6_full_loopback_test_long_enough"
resp = client.post(
"/v1/mcp/oauth/authorize",
data={
"redirect_uri": "http://[0:0:0:0:0:0:0:1]:3000/cb",
"code_challenge": _make_challenge(verifier),
"code_challenge_method": "S256",
"state": "s",
"server_id": "sid",
"api_key": "k",
},
follow_redirects=False,
)
assert resp.status_code == 302
@pytest.mark.asyncio
async def test_token_endpoint_missing_master_key_preserves_code_and_db():
"""If master_key is unset, /token must reject BEFORE consuming the
code or writing the credential — otherwise a misconfigured deploy
burns the code and persists the key with no way for the user to
retrieve a session token without restarting the flow."""
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
byok_token,
)
verifier = "verifier_for_missing_master_key_test_long_!"
challenge = _make_challenge(verifier)
code = str(uuid.uuid4())
_byok_auth_codes[code] = {
"api_key": "k",
"server_id": "sid",
"code_challenge": challenge,
"redirect_uri": "http://127.0.0.1:3000/cb",
"client_id": "",
"user_id": "u",
"expires_at": time.time() + 60,
}
mock_store = AsyncMock()
with (
patch(
"litellm.proxy._experimental.mcp_server.byok_oauth_endpoints.store_user_credential",
mock_store,
),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.master_key", None),
):
result = await byok_token(
request=MagicMock(),
grant_type="authorization_code",
code=code,
redirect_uri="http://127.0.0.1:3000/cb",
code_verifier=verifier,
client_id="",
)
assert result.status_code == 500
assert json.loads(result.body) == {"error": "server_error"}
# Code still present — user can retry once master_key is configured.
assert code in _byok_auth_codes
# Credential never written — no inconsistent DB state.
mock_store.assert_not_awaited()
def test_authorize_post_rejects_cookie_without_exp(unauthenticated_client):
"""Defense-in-depth: UI session cookies must carry an ``exp`` claim
so a leaked cookie has a bounded lifetime. A master-key-signed JWT
without ``exp`` is rejected at decode time (PyJWT
``options={"require": ["exp"]}``)."""
import jwt as _jwt
with patch("litellm.proxy.proxy_server.master_key", "real-master-key"):
no_exp = _jwt.encode(
{"user_id": "u", "login_method": "sso"},
"real-master-key",
algorithm="HS256",
)
resp = _authorize_post_with_cookie(unauthenticated_client, no_exp, api_key="k")
assert resp.status_code == 401
def test_authorize_post_rejects_expired_cookie(unauthenticated_client):
"""An expired UI session cookie is rejected, not accepted as valid."""
import jwt as _jwt
with patch("litellm.proxy.proxy_server.master_key", "real-master-key"):
expired = _jwt.encode(
{
"user_id": "u",
"login_method": "sso",
"exp": int(time.time()) - 60, # 1 minute ago
},
"real-master-key",
algorithm="HS256",
)
resp = _authorize_post_with_cookie(unauthenticated_client, expired, api_key="k")
assert resp.status_code == 401
@pytest.mark.asyncio
async def test_token_endpoint_accepts_oauth21_client_omitting_redirect_uri():
"""OAuth 2.1 draft-15 §4.1.3 dropped the redirect_uri requirement at
the token endpoint. A strict OAuth 2.1 client will omit the value —
LiteLLM must accept that and rely on PKCE + client_id binding for
the security role redirect_uri played under RFC 6749.
Enforcement still fires when the client DOES submit a value that
disagrees with the record (see test_token_endpoint_rejects_redirect_uri_mismatch).
"""
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
byok_token,
)
verifier = "verifier_for_oauth21_no_redirect_uri_omit_ok!"
challenge = _make_challenge(verifier)
code = str(uuid.uuid4())
_byok_auth_codes[code] = {
"api_key": "k",
"server_id": "sid",
"code_challenge": challenge,
"redirect_uri": "http://127.0.0.1:3000/cb",
"client_id": "claude-desktop",
"user_id": "u",
"expires_at": time.time() + 60,
}
with (
patch(
"litellm.proxy._experimental.mcp_server.byok_oauth_endpoints.store_user_credential",
new=AsyncMock(),
),
patch("litellm.proxy.proxy_server.master_key", "test-master"),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
):
result = await byok_token(
request=MagicMock(),
grant_type="authorization_code",
code=code,
redirect_uri="", # OAuth 2.1 client omits it
code_verifier=verifier,
client_id="claude-desktop",
)
assert result.status_code == 200
def test_validate_loopback_redirect_uri_rejects_fragment():
from litellm.proxy._experimental.mcp_server.oauth_utils import (
validate_loopback_redirect_uri,
)
with pytest.raises(HTTPException) as exc:
validate_loopback_redirect_uri("http://127.0.0.1:3000/cb#code=1")
assert exc.value.status_code == 400
def test_validate_loopback_redirect_uri_rejects_malformed_cleanly():
"""Malformed / unparseable URIs should surface as 400 invalid_request,
not a 500 from an unhandled exception inside ip_address()."""
from litellm.proxy._experimental.mcp_server.oauth_utils import (
validate_loopback_redirect_uri,
)
# Netloc that parses but whose host is neither "localhost" nor a valid IP.
with pytest.raises(HTTPException) as exc:
validate_loopback_redirect_uri("http://[not-an-ip]/cb")
assert exc.value.status_code == 400

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

@ -76,7 +76,7 @@ async def test_authorize_endpoint_includes_response_type():
request=mock_request,
client_id="test_client_id",
mcp_server_name="test_oauth",
redirect_uri="https://client.example.com/callback",
redirect_uri="http://127.0.0.1:60108/callback",
state="test_state",
)
@ -139,7 +139,7 @@ async def test_authorize_endpoint_preserves_existing_query_params():
request=mock_request,
client_id="test_client_id",
mcp_server_name="test_oauth",
redirect_uri="https://client.example.com/callback",
redirect_uri="http://127.0.0.1:60108/callback",
state="test_state",
)
@ -558,7 +558,7 @@ async def test_authorize_endpoint_respects_x_forwarded_proto():
request=mock_request,
client_id="test_client_id",
mcp_server_name="test_oauth",
redirect_uri="https://client.example.com/callback",
redirect_uri="http://127.0.0.1:60108/callback",
state="test_state",
)
@ -855,7 +855,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host():
request=mock_request,
client_id="test_client_id",
mcp_server_name="test_oauth",
redirect_uri="https://client.example.com/callback",
redirect_uri="http://127.0.0.1:60108/callback",
state="test_state",
)
@ -1817,3 +1817,191 @@ async def test_token_endpoint_authorization_code_missing_code():
)
assert exc_info.value.status_code == 400
assert "code is required" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_authorize_endpoint_rejects_non_loopback_redirect_uri():
"""VERIA-57 root cause B regression. The client-supplied redirect_uri
is encrypted into the OAuth state and decoded on /callback to 302 the
user back. A non-loopback value is an open-redirect + code-theft
primitive — reject with 400 before encoding anything into state."""
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
authorize,
)
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
global_mcp_server_manager.registry.clear()
oauth2_server = MCPServer(
server_id="test_oauth_server",
name="test_oauth",
server_name="test_oauth",
alias="test_oauth",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
client_secret="cs",
authorization_url="https://provider.com/oauth/authorize",
token_url="https://provider.com/oauth/token",
)
global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server
mock_request = MagicMock(spec=Request)
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
with pytest.raises(HTTPException) as exc_info:
await authorize(
request=mock_request,
client_id="cid",
mcp_server_name="test_oauth",
redirect_uri="https://attacker.example.com/cb",
state="s",
)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_authorize_endpoint_accepts_ipv4_loopback_range_and_ipv6_full_form():
"""RFC 8252 §7.3 + RFC 4291: full 127.0.0.0/8 and full-form IPv6
loopback must be accepted — string match on ``127.0.0.1`` alone
would miss ``127.0.0.2`` and ``0:0:0:0:0:0:0:1``."""
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
authorize,
)
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
global_mcp_server_manager.registry.clear()
oauth2_server = MCPServer(
server_id="test_oauth_server",
name="test_oauth",
server_name="test_oauth",
alias="test_oauth",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
client_secret="cs",
authorization_url="https://provider.com/oauth/authorize",
token_url="https://provider.com/oauth/token",
)
global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server
mock_request = MagicMock(spec=Request)
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
for uri in (
"http://127.0.0.2:3000/cb",
"http://[0:0:0:0:0:0:0:1]:3000/cb",
"http://localhost:3000/cb",
):
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper"
) as mock_encrypt:
mock_encrypt.return_value = "mocked_encrypted_state"
response = await authorize(
request=mock_request,
client_id="cid",
mcp_server_name="test_oauth",
redirect_uri=uri,
state="s",
)
assert response.status_code == 307, f"{uri} should be accepted"
@pytest.mark.asyncio
async def test_callback_revalidates_loopback_on_decoded_base_url():
"""VERIA-57 root cause B defense-in-depth: an encrypted state minted
before the /authorize validation was added has no expiry and stays
valid. /callback must re-validate the decoded base_url so those
stale states can't be used as an open-redirect + code-theft
primitive."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
callback,
)
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash"
) as mock_decode:
mock_decode.return_value = {
"base_url": "https://attacker.example.com/cb",
"original_state": "s",
"code_challenge": None,
"code_challenge_method": None,
"client_redirect_uri": "https://attacker.example.com/cb",
}
with pytest.raises(HTTPException) as exc_info:
await callback(code="stolen_code", state="encrypted_stale_state")
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_token_endpoint_sets_no_store_cache_control():
"""RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: the token response
contains an access token (and possibly a refresh token) — it MUST
NOT be cached by intermediaries or the client."""
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
exchange_token_with_server,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="t",
name="t",
server_name="t",
alias="t",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
client_secret="cs",
authorization_url="https://provider.com/oauth/authorize",
token_url="https://provider.com/oauth/token",
)
mock_request = MagicMock(spec=Request)
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
fake_http_response = MagicMock()
fake_http_response.json.return_value = {
"access_token": "tok",
"token_type": "Bearer",
"expires_in": 3600,
}
fake_http_response.raise_for_status = MagicMock()
fake_http_client = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
return_value=fake_http_client,
):
response = await exchange_token_with_server(
request=mock_request,
mcp_server=server,
grant_type="authorization_code",
code="c",
redirect_uri="http://127.0.0.1:3000/cb",
client_id="cid",
client_secret=None,
code_verifier=None,
)
assert response.headers["cache-control"] == "no-store"
assert response.headers["pragma"] == "no-cache"