mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(mcp): persist DCR client_id so interactive OAuth token refresh works (#31912)
* fix(mcp): persist DCR client_id so interactive OAuth token refresh works
Interactive authorization_code MCP servers register an OAuth client via Dynamic
Client Registration (RFC 7591) during the authorize flow, but the minted
client_id and the discovered token_url were returned to the caller and never
written to the server row. The autonomous refresh_token grant reads client_id,
client_secret and token_url off the server, so an expired access token could not
be refreshed; the user was bounced back to re-authorize and tools/list returned
zero tools
Persist the DCR client_id (plus client_secret and token_endpoint_auth_method when
the registration returns them) and the discovered token_url onto the server row,
reusing the encrypt_credentials write that client_credentials and token exchange
already use, then refresh the in-memory registry so the value is live at refresh
time. Both the v1 refresher and the v2 AuthorizationCodeRefresher read those same
fields, so egress needs no change
* fix: reuse persisted MCP DCR clients
* fix: reuse persisted MCP DCR clients
---------
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
(cherry picked from commit 15ff389eb4)
This commit is contained in:
parent
13db8cb140
commit
fe9fb0b81a
4 changed files with 778 additions and 3 deletions
|
|
@ -2,12 +2,13 @@ import asyncio
|
|||
import html as _html
|
||||
import json
|
||||
import time
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple
|
||||
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Form, HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -30,9 +31,12 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
|
||||
from litellm.proxy.utils import get_server_root_path
|
||||
from litellm.types.mcp import MCPAuth
|
||||
from litellm.types.mcp import MCPAuth, MCPCredentials
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import LiteLLM_MCPServerTable
|
||||
|
||||
# TTL cache for upstream OAuth metadata fetched from pass-through MCP servers.
|
||||
# Keeps us from hammering the upstream IdP on each discovery request.
|
||||
# Keyed by (server_id, resource_url) → (expires_at_epoch, payload).
|
||||
|
|
@ -509,6 +513,178 @@ async def exchange_token_with_server(
|
|||
return JSONResponse(result, headers=TOKEN_NO_CACHE_HEADERS)
|
||||
|
||||
|
||||
class _DcrClientRegistration(BaseModel):
|
||||
"""RFC 7591 dynamic client registration response, narrowed to the fields the gateway
|
||||
must persist to authenticate later token-endpoint calls. Extra members are ignored."""
|
||||
|
||||
client_id: str
|
||||
client_secret: Optional[str] = None
|
||||
token_endpoint_auth_method: Optional[str] = None
|
||||
|
||||
|
||||
class _PersistedDcrCredentials(BaseModel):
|
||||
client_id: Optional[str] = None
|
||||
client_secret: Optional[str] = None
|
||||
token_endpoint_auth_method: Optional[str] = None
|
||||
|
||||
|
||||
def _get_persisted_dcr_credentials(credentials: object) -> Optional[_PersistedDcrCredentials]:
|
||||
if not credentials:
|
||||
return None
|
||||
try:
|
||||
return (
|
||||
_PersistedDcrCredentials.model_validate_json(credentials)
|
||||
if isinstance(credentials, str)
|
||||
else _PersistedDcrCredentials.model_validate(credentials)
|
||||
)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _decrypt_persisted_dcr_credential(value: Optional[str], key: str) -> Optional[str]:
|
||||
if value is None:
|
||||
return None
|
||||
return decrypt_value_helper(
|
||||
value=value,
|
||||
key=key,
|
||||
exception_type="debug",
|
||||
return_original_value=True,
|
||||
)
|
||||
|
||||
|
||||
def _apply_persisted_dcr_credentials(mcp_server: MCPServer, credentials: _PersistedDcrCredentials) -> bool:
|
||||
client_id = _decrypt_persisted_dcr_credential(credentials.client_id, "client_id")
|
||||
if not client_id:
|
||||
return False
|
||||
mcp_server.client_id = client_id
|
||||
mcp_server.client_secret = _decrypt_persisted_dcr_credential(credentials.client_secret, "client_secret")
|
||||
mcp_server.token_endpoint_auth_method = credentials.token_endpoint_auth_method
|
||||
return True
|
||||
|
||||
|
||||
async def _get_persisted_mcp_server_with_dcr_client_id(
|
||||
mcp_server: MCPServer,
|
||||
) -> Optional[tuple["LiteLLM_MCPServerTable", _PersistedDcrCredentials]]:
|
||||
from litellm.proxy._experimental.mcp_server.db import get_mcp_server # noqa: PLC0415
|
||||
from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415
|
||||
|
||||
try:
|
||||
prisma_client = get_prisma_client_or_throw("Database not connected. Cannot read MCP OAuth client registration.")
|
||||
persisted_mcp_server = await get_mcp_server(
|
||||
prisma_client=prisma_client,
|
||||
server_id=mcp_server.server_id,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
verbose_logger.debug(
|
||||
"register_client_with_server: failed to read persisted DCR client registration for server_id=%s: %s",
|
||||
mcp_server.server_id,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
if persisted_mcp_server is None:
|
||||
return None
|
||||
|
||||
credentials = _get_persisted_dcr_credentials(persisted_mcp_server.credentials)
|
||||
if credentials is None or not credentials.client_id:
|
||||
return None
|
||||
|
||||
return persisted_mcp_server, credentials
|
||||
|
||||
|
||||
async def _reuse_persisted_dcr_client_if_available(mcp_server: MCPServer) -> bool:
|
||||
persisted = await _get_persisted_mcp_server_with_dcr_client_id(mcp_server)
|
||||
if persisted is None:
|
||||
return False
|
||||
persisted_mcp_server, credentials = persisted
|
||||
if not _apply_persisted_dcr_credentials(mcp_server, credentials):
|
||||
return False
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
try:
|
||||
await global_mcp_server_manager.update_server(persisted_mcp_server)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
verbose_logger.warning(
|
||||
"register_client_with_server: failed to refresh persisted DCR client registration for server_id=%s: %s",
|
||||
mcp_server.server_id,
|
||||
exc,
|
||||
)
|
||||
return bool(mcp_server.client_id)
|
||||
|
||||
|
||||
DcrRegistrationPersistenceResult = Literal["persisted", "reused", "failed"]
|
||||
|
||||
|
||||
async def _persist_dcr_client_registration(
|
||||
mcp_server: MCPServer, registration_response: object
|
||||
) -> DcrRegistrationPersistenceResult:
|
||||
"""Persist the dynamically registered OAuth client (RFC 7591) onto the MCP server row.
|
||||
|
||||
The interactive authorization_code flow mints a ``client_id`` via Dynamic Client
|
||||
Registration that discovery cannot re-derive; without persisting it the autonomous
|
||||
``refresh_token`` grant has no client identity, so an expired access token forces a
|
||||
full re-authorization instead of a silent refresh. Mirrors the ``encrypt_credentials``
|
||||
write that ``client_credentials`` and token exchange already use. Failures are logged,
|
||||
never raised: registration still returns to the caller even when persistence fails.
|
||||
"""
|
||||
try:
|
||||
registration = _DcrClientRegistration.model_validate(registration_response)
|
||||
except ValidationError as exc:
|
||||
verbose_logger.warning(
|
||||
"register_client_with_server: DCR response has no usable client_id for server_id=%s; "
|
||||
"client registration not persisted (%s)",
|
||||
mcp_server.server_id,
|
||||
exc,
|
||||
)
|
||||
return "failed"
|
||||
|
||||
if await _reuse_persisted_dcr_client_if_available(mcp_server):
|
||||
return "reused"
|
||||
|
||||
credentials: MCPCredentials = {
|
||||
"client_id": registration.client_id,
|
||||
**({"client_secret": registration.client_secret} if registration.client_secret is not None else {}),
|
||||
**(
|
||||
{"token_endpoint_auth_method": "client_secret_basic"}
|
||||
if registration.token_endpoint_auth_method == "client_secret_basic"
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import update_mcp_server # noqa: PLC0415
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.proxy._types import UpdateMCPServerRequest # noqa: PLC0415
|
||||
from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415
|
||||
|
||||
try:
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Cannot persist MCP OAuth client registration."
|
||||
)
|
||||
updated_row = await update_mcp_server(
|
||||
prisma_client=prisma_client,
|
||||
data=UpdateMCPServerRequest(
|
||||
server_id=mcp_server.server_id,
|
||||
credentials=credentials,
|
||||
**({"token_url": mcp_server.token_url} if mcp_server.token_url else {}),
|
||||
),
|
||||
touched_by="mcp_oauth_dcr",
|
||||
)
|
||||
await global_mcp_server_manager.update_server(updated_row)
|
||||
return "persisted"
|
||||
except Exception as exc: # noqa: BLE001
|
||||
verbose_logger.warning(
|
||||
"register_client_with_server: failed to persist DCR client registration for server_id=%s: %s",
|
||||
mcp_server.server_id,
|
||||
exc,
|
||||
)
|
||||
return "failed"
|
||||
|
||||
|
||||
async def register_client_with_server(
|
||||
request: Request,
|
||||
mcp_server: MCPServer,
|
||||
|
|
@ -517,6 +693,7 @@ async def register_client_with_server(
|
|||
response_types: Optional[list],
|
||||
token_endpoint_auth_method: Optional[str],
|
||||
fallback_client_id: Optional[str] = None,
|
||||
persist_credentials: bool = False,
|
||||
):
|
||||
request_base_url = get_request_base_url(request)
|
||||
dummy_return = {
|
||||
|
|
@ -525,7 +702,10 @@ async def register_client_with_server(
|
|||
"redirect_uris": [f"{request_base_url}/callback"],
|
||||
}
|
||||
|
||||
if mcp_server.client_id and mcp_server.client_secret:
|
||||
if mcp_server.client_id:
|
||||
return dummy_return
|
||||
|
||||
if await _reuse_persisted_dcr_client_if_available(mcp_server):
|
||||
return dummy_return
|
||||
|
||||
if mcp_server.authorization_url is None:
|
||||
|
|
@ -561,6 +741,11 @@ async def register_client_with_server(
|
|||
|
||||
token_response = response.json()
|
||||
|
||||
if persist_credentials:
|
||||
persistence_result = await _persist_dcr_client_registration(mcp_server, token_response)
|
||||
if persistence_result == "reused":
|
||||
return dummy_return
|
||||
|
||||
return JSONResponse(token_response)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1714,6 +1714,7 @@ if MCP_AVAILABLE:
|
|||
response_types=data.get("response_types", []),
|
||||
token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""),
|
||||
fallback_client_id=server_id,
|
||||
persist_credentials=_user_is_full_admin(user_api_key_dict),
|
||||
)
|
||||
|
||||
@router.delete(
|
||||
|
|
|
|||
|
|
@ -534,6 +534,548 @@ async def test_register_client_remote_registration_success():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_client_persists_dcr_client_identity():
|
||||
"""A dynamic client registration (RFC 7591) must persist the issued client_id /
|
||||
client_secret / token_endpoint_auth_method and the token_url onto the server row so
|
||||
autonomous refresh can authenticate as the registered client. Without persistence the
|
||||
minted client_id is discarded and the refresh_token grant has no client identity."""
|
||||
try:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client_with_server,
|
||||
)
|
||||
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")
|
||||
|
||||
oauth2_server = MCPServer(
|
||||
server_id="remote_server",
|
||||
name="remote_server",
|
||||
server_name="remote_server",
|
||||
alias="remote_server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
client_id=None,
|
||||
client_secret=None,
|
||||
authorization_url="https://provider.example/oauth/authorize",
|
||||
token_url="https://provider.example/oauth/token",
|
||||
registration_url="https://provider.example/oauth/register",
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.base_url = "https://proxy.litellm.example/"
|
||||
mock_request.headers = {}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"client_id": "generated-client",
|
||||
"client_secret": "generated-secret",
|
||||
"token_endpoint_auth_method": "client_secret_basic",
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_async_client = MagicMock()
|
||||
mock_async_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
mock_update = AsyncMock(return_value=MagicMock())
|
||||
mock_update_server = AsyncMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
|
||||
return_value=mock_async_client,
|
||||
),
|
||||
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()),
|
||||
patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update),
|
||||
patch.object(global_mcp_server_manager, "update_server", new=mock_update_server),
|
||||
):
|
||||
response = await register_client_with_server(
|
||||
request=mock_request,
|
||||
mcp_server=oauth2_server,
|
||||
client_name="Litellm Proxy",
|
||||
grant_types=["authorization_code", "refresh_token"],
|
||||
response_types=["code"],
|
||||
token_endpoint_auth_method="client_secret_basic",
|
||||
persist_credentials=True,
|
||||
)
|
||||
|
||||
import json
|
||||
|
||||
assert response.status_code == 200
|
||||
assert json.loads(response.body.decode("utf-8")) == mock_response.json.return_value
|
||||
|
||||
mock_update.assert_called_once()
|
||||
update_data = mock_update.call_args.kwargs["data"]
|
||||
assert update_data.server_id == "remote_server"
|
||||
assert update_data.token_url == "https://provider.example/oauth/token"
|
||||
assert update_data.credentials["client_id"] == "generated-client"
|
||||
assert update_data.credentials["client_secret"] == "generated-secret"
|
||||
assert update_data.credentials["token_endpoint_auth_method"] == "client_secret_basic"
|
||||
|
||||
mock_update_server.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_client_does_not_clobber_token_url_when_absent():
|
||||
"""When the in-memory server has no token_url, the DCR persist must omit it from the
|
||||
partial update rather than passing None, so exclude_unset leaves the token_url column
|
||||
untouched instead of overwriting an existing value with NULL."""
|
||||
try:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client_with_server,
|
||||
)
|
||||
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")
|
||||
|
||||
oauth2_server = MCPServer(
|
||||
server_id="remote_server",
|
||||
name="remote_server",
|
||||
server_name="remote_server",
|
||||
alias="remote_server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
client_id=None,
|
||||
client_secret=None,
|
||||
authorization_url="https://provider.example/oauth/authorize",
|
||||
token_url=None,
|
||||
registration_url="https://provider.example/oauth/register",
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.base_url = "https://proxy.litellm.example/"
|
||||
mock_request.headers = {}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"client_id": "generated-client"}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_async_client = MagicMock()
|
||||
mock_async_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
mock_update = AsyncMock(return_value=MagicMock())
|
||||
mock_update_server = AsyncMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
|
||||
return_value=mock_async_client,
|
||||
),
|
||||
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()),
|
||||
patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update),
|
||||
patch.object(global_mcp_server_manager, "update_server", new=mock_update_server),
|
||||
):
|
||||
await register_client_with_server(
|
||||
request=mock_request,
|
||||
mcp_server=oauth2_server,
|
||||
client_name="Litellm Proxy",
|
||||
grant_types=["authorization_code", "refresh_token"],
|
||||
response_types=["code"],
|
||||
token_endpoint_auth_method="none",
|
||||
persist_credentials=True,
|
||||
)
|
||||
|
||||
mock_update.assert_called_once()
|
||||
update_data = mock_update.call_args.kwargs["data"]
|
||||
assert update_data.credentials["client_id"] == "generated-client"
|
||||
assert "token_url" not in update_data.model_fields_set
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_client_reuses_persisted_client_id_for_non_admin_when_registry_is_stale():
|
||||
try:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client_with_server,
|
||||
)
|
||||
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")
|
||||
|
||||
oauth2_server = MCPServer(
|
||||
server_id="remote_server",
|
||||
name="remote_server",
|
||||
server_name="remote_server",
|
||||
alias="remote_server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
client_id=None,
|
||||
client_secret=None,
|
||||
authorization_url="https://provider.example/oauth/authorize",
|
||||
token_url="https://provider.example/oauth/token",
|
||||
registration_url="https://provider.example/oauth/register",
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.base_url = "https://proxy.litellm.example/"
|
||||
mock_request.headers = {}
|
||||
|
||||
persisted_server = MagicMock()
|
||||
persisted_server.credentials = {"client_id": "persisted-client"}
|
||||
mock_get_mcp_server = AsyncMock(return_value=persisted_server)
|
||||
mock_update_mcp_server = AsyncMock()
|
||||
mock_update_server = AsyncMock()
|
||||
mock_async_client = MagicMock()
|
||||
mock_async_client.post = AsyncMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
|
||||
return_value=mock_async_client,
|
||||
),
|
||||
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.get_mcp_server",
|
||||
new=mock_get_mcp_server,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.update_mcp_server",
|
||||
new=mock_update_mcp_server,
|
||||
),
|
||||
patch.object(global_mcp_server_manager, "update_server", new=mock_update_server),
|
||||
):
|
||||
response = await register_client_with_server(
|
||||
request=mock_request,
|
||||
mcp_server=oauth2_server,
|
||||
client_name="Litellm Proxy",
|
||||
grant_types=["authorization_code", "refresh_token"],
|
||||
response_types=["code"],
|
||||
token_endpoint_auth_method="none",
|
||||
persist_credentials=False,
|
||||
)
|
||||
|
||||
assert response["client_id"] == "remote_server"
|
||||
assert oauth2_server.client_id == "persisted-client"
|
||||
mock_async_client.post.assert_not_called()
|
||||
mock_update_mcp_server.assert_not_called()
|
||||
mock_update_server.assert_called_once_with(persisted_server)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_client_reuse_refreshes_request_server_when_manager_update_fails():
|
||||
try:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client_with_server,
|
||||
)
|
||||
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")
|
||||
|
||||
oauth2_server = MCPServer(
|
||||
server_id="remote_server",
|
||||
name="remote_server",
|
||||
server_name="remote_server",
|
||||
alias="remote_server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
client_id=None,
|
||||
client_secret=None,
|
||||
authorization_url="https://provider.example/oauth/authorize",
|
||||
token_url="https://provider.example/oauth/token",
|
||||
registration_url="https://provider.example/oauth/register",
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.base_url = "https://proxy.litellm.example/"
|
||||
mock_request.headers = {}
|
||||
|
||||
persisted_server = MagicMock()
|
||||
persisted_server.credentials = {
|
||||
"client_id": "persisted-client",
|
||||
"client_secret": "persisted-secret",
|
||||
"token_endpoint_auth_method": "client_secret_basic",
|
||||
}
|
||||
mock_get_mcp_server = AsyncMock(return_value=persisted_server)
|
||||
mock_async_client = MagicMock()
|
||||
mock_async_client.post = AsyncMock()
|
||||
mock_update_server = AsyncMock(side_effect=RuntimeError("registry update failed"))
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
|
||||
return_value=mock_async_client,
|
||||
),
|
||||
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.get_mcp_server",
|
||||
new=mock_get_mcp_server,
|
||||
),
|
||||
patch.object(global_mcp_server_manager, "update_server", new=mock_update_server),
|
||||
):
|
||||
response = await register_client_with_server(
|
||||
request=mock_request,
|
||||
mcp_server=oauth2_server,
|
||||
client_name="Litellm Proxy",
|
||||
grant_types=["authorization_code", "refresh_token"],
|
||||
response_types=["code"],
|
||||
token_endpoint_auth_method="none",
|
||||
persist_credentials=False,
|
||||
)
|
||||
|
||||
assert response["client_id"] == "remote_server"
|
||||
assert oauth2_server.client_id == "persisted-client"
|
||||
assert oauth2_server.client_secret == "persisted-secret"
|
||||
assert oauth2_server.token_endpoint_auth_method == "client_secret_basic"
|
||||
mock_async_client.post.assert_not_called()
|
||||
mock_update_server.assert_called_once_with(persisted_server)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_client_returns_reused_client_when_concurrent_persist_wins():
|
||||
try:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client_with_server,
|
||||
)
|
||||
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")
|
||||
|
||||
oauth2_server = MCPServer(
|
||||
server_id="remote_server",
|
||||
name="remote_server",
|
||||
server_name="remote_server",
|
||||
alias="remote_server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
client_id=None,
|
||||
client_secret=None,
|
||||
authorization_url="https://provider.example/oauth/authorize",
|
||||
token_url="https://provider.example/oauth/token",
|
||||
registration_url="https://provider.example/oauth/register",
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.base_url = "https://proxy.litellm.example/"
|
||||
mock_request.headers = {}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"client_id": "generated-client"}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_async_client = MagicMock()
|
||||
mock_async_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
persisted_server = MagicMock()
|
||||
persisted_server.credentials = {"client_id": "persisted-client"}
|
||||
mock_get_mcp_server = AsyncMock(side_effect=[None, persisted_server])
|
||||
mock_update_mcp_server = AsyncMock()
|
||||
mock_update_server = AsyncMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
|
||||
return_value=mock_async_client,
|
||||
),
|
||||
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.get_mcp_server",
|
||||
new=mock_get_mcp_server,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.update_mcp_server",
|
||||
new=mock_update_mcp_server,
|
||||
),
|
||||
patch.object(global_mcp_server_manager, "update_server", new=mock_update_server),
|
||||
):
|
||||
response = await register_client_with_server(
|
||||
request=mock_request,
|
||||
mcp_server=oauth2_server,
|
||||
client_name="Litellm Proxy",
|
||||
grant_types=["authorization_code", "refresh_token"],
|
||||
response_types=["code"],
|
||||
token_endpoint_auth_method="none",
|
||||
persist_credentials=True,
|
||||
)
|
||||
|
||||
assert response["client_id"] == "remote_server"
|
||||
assert oauth2_server.client_id == "persisted-client"
|
||||
mock_async_client.post.assert_called_once()
|
||||
mock_update_mcp_server.assert_not_called()
|
||||
mock_update_server.assert_called_once_with(persisted_server)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_client_reuses_existing_client_id_without_re_dcr():
|
||||
"""A server that already has a client_id (admin-configured or previously DCR'd) must be
|
||||
reused, not re-registered, even without a client_secret. A client_id is one-per-application
|
||||
in OAuth and shared across users; re-minting per authorize would orphan other users' refresh
|
||||
tokens by overwriting the server's client_id."""
|
||||
try:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client,
|
||||
)
|
||||
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")
|
||||
|
||||
global_mcp_server_manager.registry.clear()
|
||||
oauth2_server = MCPServer(
|
||||
server_id="remote_server",
|
||||
name="remote_server",
|
||||
server_name="remote_server",
|
||||
alias="remote_server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
client_id="existing-shared-client",
|
||||
client_secret=None,
|
||||
authorization_url="https://provider.example/oauth/authorize",
|
||||
token_url="https://provider.example/oauth/token",
|
||||
registration_url="https://provider.example/oauth/register",
|
||||
)
|
||||
global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.base_url = "https://proxy.litellm.example/"
|
||||
mock_request.headers = {}
|
||||
|
||||
request_payload = {
|
||||
"client_name": "Litellm Proxy",
|
||||
"grant_types": ["authorization_code", "refresh_token"],
|
||||
"response_types": ["code"],
|
||||
"token_endpoint_auth_method": "none",
|
||||
}
|
||||
|
||||
mock_async_client = MagicMock()
|
||||
mock_async_client.post = AsyncMock()
|
||||
|
||||
try:
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body",
|
||||
new=AsyncMock(return_value=request_payload),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
|
||||
return_value=mock_async_client,
|
||||
),
|
||||
):
|
||||
response = await register_client(request=mock_request, mcp_server_name=oauth2_server.server_name)
|
||||
finally:
|
||||
global_mcp_server_manager.registry.clear()
|
||||
|
||||
mock_async_client.post.assert_not_called()
|
||||
body = response if isinstance(response, dict) else json.loads(response.body.decode("utf-8"))
|
||||
assert body["client_secret"] == "dummy"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_register_route_does_not_persist_client_credentials():
|
||||
"""The unauthenticated root /register route must not persist the DCR result onto the
|
||||
server row; only the authenticated management path passes persist_credentials=True. An
|
||||
external caller could otherwise bind a caller-controlled client (and leak its secret) to
|
||||
a server that has no client yet."""
|
||||
try:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client,
|
||||
)
|
||||
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")
|
||||
|
||||
global_mcp_server_manager.registry.clear()
|
||||
oauth2_server = MCPServer(
|
||||
server_id="remote_server",
|
||||
name="remote_server",
|
||||
server_name="remote_server",
|
||||
alias="remote_server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
client_id=None,
|
||||
client_secret=None,
|
||||
authorization_url="https://provider.example/oauth/authorize",
|
||||
token_url="https://provider.example/oauth/token",
|
||||
registration_url="https://provider.example/oauth/register",
|
||||
)
|
||||
global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.base_url = "https://proxy.litellm.example/"
|
||||
mock_request.headers = {}
|
||||
|
||||
request_payload = {
|
||||
"client_name": "attacker",
|
||||
"grant_types": ["authorization_code"],
|
||||
"response_types": ["code"],
|
||||
"token_endpoint_auth_method": "none",
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"client_id": "attacker-client",
|
||||
"client_secret": "attacker-secret",
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_async_client = MagicMock()
|
||||
mock_async_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
mock_update = AsyncMock(return_value=MagicMock())
|
||||
|
||||
try:
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body",
|
||||
new=AsyncMock(return_value=request_payload),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
|
||||
return_value=mock_async_client,
|
||||
),
|
||||
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()),
|
||||
patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()),
|
||||
patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update),
|
||||
):
|
||||
await register_client(request=mock_request, mcp_server_name=oauth2_server.server_name)
|
||||
finally:
|
||||
global_mcp_server_manager.registry.clear()
|
||||
|
||||
mock_update.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("trust_xff")
|
||||
async def test_authorize_endpoint_respects_x_forwarded_proto():
|
||||
|
|
|
|||
|
|
@ -2264,8 +2264,55 @@ class TestTemporaryMCPSessionEndpoints:
|
|||
response_types=["code"],
|
||||
token_endpoint_auth_method="client_secret_basic",
|
||||
fallback_client_id="server-1",
|
||||
persist_credentials=True,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_register_does_not_persist_for_non_admin(self):
|
||||
"""A non-admin caller (who may have access to a real server) must not persist the DCR
|
||||
result onto the shared server row. register_client_with_server is invoked with
|
||||
persist_credentials=False, so user-side registration returns the DCR response without
|
||||
writing shared client credentials. Only a full PROXY_ADMIN establishes the shared client."""
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
mcp_register,
|
||||
)
|
||||
|
||||
request = MagicMock()
|
||||
server = generate_mock_mcp_server_config_record(server_id="server-1")
|
||||
register_response = {"client_id": "generated"}
|
||||
request_body = {
|
||||
"client_name": "LiteLLM",
|
||||
"grant_types": ["authorization_code"],
|
||||
"response_types": ["code"],
|
||||
"token_endpoint_auth_method": "client_secret_basic",
|
||||
}
|
||||
non_admin_auth = generate_mock_user_api_key_auth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404",
|
||||
return_value=server,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body",
|
||||
AsyncMock(return_value=request_body),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.register_client_with_server",
|
||||
AsyncMock(return_value=register_response),
|
||||
) as register_mock,
|
||||
):
|
||||
result = await mcp_register(
|
||||
request=request,
|
||||
server_id="server-1",
|
||||
user_api_key_dict=non_admin_auth,
|
||||
)
|
||||
|
||||
assert result is register_response
|
||||
assert register_mock.await_args.kwargs["persist_credentials"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_temporary_mcp_server_falls_back_to_redis(self):
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue