scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-18 04:44:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

fix(mcp)!: register the gateway's own DCR client for bridge servers instead of relaying the caller's loopback redirect (#46940)

Browse source 
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-10-17 12:14:28 -07:00 • committed by GitHub
parent 9e9dccb9d4
commit 7049a5ff4d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 481 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -862,12 +862,17 @@ def _raise_unless_oauth2_discovery_server(
def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool:
"""True when a DCR-bridge server relays client registration to the upstream authorization
server instead of short-circuiting to an admin-configured OAuth client. In the relay arm the
server instead of short-circuiting to the server's own OAuth client. In the relay arm the
upstream holds each client's own registration, so the authorize and token relays pass the
client's ``client_id`` and ``redirect_uri`` through verbatim and the authorization code
returns directly to the client's redirect URI without transiting the gateway. Gateway-side
redirect trust and the ``/callback`` state relay therefore only apply to the short-circuit
arm, where the upstream only knows the gateway's own callback."""
arm, where the upstream only knows the gateway's own callback.
Authorize and token decide per request from the server's current client: once the gateway
holds its shared client (registered at ``/register`` under ``PROXY_BASE_URL``, or entered by an
admin) every caller runs the gateway-callback arm, while a server that still has none keeps
relaying, so clients registered before the shared client existed keep their own binding."""
return bool(
mcp_server.is_dcr_bridge
and mcp_server.effective_registration_url
@ -880,6 +885,27 @@ def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool:
)
def _register_relays_bridge_client(mcp_server: MCPServer) -> bool:
"""``/register`` relays the caller's own registration only when the gateway cannot keep a shared
client: with ``PROXY_BASE_URL`` set it registers one client bound to its own ``/callback`` and
keeps it as the server's client, so an upstream that restricts redirect domains (Atlassian)
never sees a laptop loopback redirect. Without a database the shared client would live only in
this process, so a database-free gateway keeps relaying as well."""
return _dcr_bridge_relays_client_registration(mcp_server) and (
configured_proxy_base_url() is None or not _shared_client_store_available()
)
def _shared_client_store_available() -> bool:
from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # avoids circular import
try:
get_prisma_client_or_throw("Database not connected. Cannot keep the shared MCP OAuth client.")
except HTTPException:
return False
return True
def _require_s256_pkce(
code_challenge: str | None,
code_challenge_method: str | None,
@ -1028,6 +1054,17 @@ async def _managed_authorization_scope(
return selected_scope, scope_user_id
async def _hydrate_bridge_shared_client(mcp_server: MCPServer, request: Request) -> None:
"""Whichever worker served ``/register`` persisted the bridge's shared gateway client; a worker
whose catalog snapshot predates that write still holds no ``client_id``, so the gateway-callback
arms load the persisted client before presenting one upstream. Callers that carry their own
client (a gateway-minted ephemeral client sealed into the state or the code) skip this, since the
upstream bound their code to that client."""
if not mcp_server.is_dcr_bridge or mcp_server.client_id:
return
await _reuse_persisted_dcr_client_if_available(mcp_server, f"{get_request_base_url(request)}/callback")
async def authorize_with_server(
request: Request,
mcp_server: MCPServer,
@ -1069,6 +1106,8 @@ async def authorize_with_server(
# calling this for its enforcement side effect, then falls through to the gateway
# /callback flow below, which reads the original code_challenge names.
bridge_challenge, bridge_method = _require_s256_pkce(code_challenge, code_challenge_method)
if ephemeral_dcr_client is None:
await _hydrate_bridge_shared_client(resolved_server, request)
# A gateway-minted ephemeral client is registered against {base}/callback, so its
# flow must run the short-circuit arm; the relay arm is only for clients that
# registered themselves through the front door and hold their own redirect binding.
@ -1233,6 +1272,7 @@ async def exchange_token_with_server(
refresh_token: str | None = None,
scope: str | None = None,
client_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None,
sealed_client: bool = False,
):
raise_if_not_oauth2(mcp_server)
if grant_type not in ("authorization_code", "refresh_token"):
@ -1282,11 +1322,14 @@ async def exchange_token_with_server(
# recovered from a sealed code) must authenticate the way its own registration was granted,
# not the way the server row is configured; callers that carry no method keep the row's method
# as before.
resolved_client_id: Final = resolved_server.client_id if resolved_server.client_id else client_id
resolved_client_secret: Final = resolved_server.client_secret if resolved_server.client_id else client_secret
if not sealed_client:
await _hydrate_bridge_shared_client(resolved_server, request)
server_client_wins: Final = bool(resolved_server.client_id) and not sealed_client
resolved_client_id: Final = resolved_server.client_id if server_client_wins else client_id
resolved_client_secret: Final = resolved_server.client_secret if server_client_wins else client_secret
resolved_auth_method: Final = (
resolved_server.token_endpoint_auth_method
if resolved_server.client_id
if server_client_wins
else (client_token_endpoint_auth_method or resolved_server.token_endpoint_auth_method)
)
try:
@ -1830,10 +1873,13 @@ async def _persist_dcr_client_registration(
empty and otherwise adopts whatever a concurrent registration persisted first.
The client-forwarded token modes (``true_passthrough`` / ``oauth_delegate``) are skipped
unconditionally: the caller holds the upstream token and the gateway must hold no OAuth
client identity for these servers. Persisting here would stamp ``oauth2_flow`` and a
``client_id`` onto a server whose mode promises the gateway stores nothing, making a
fresh pass-through server read as gateway-authorized.
unless the server is a DCR bridge: the caller holds the upstream token and the gateway holds
no OAuth client identity for these servers, so persisting would stamp ``oauth2_flow`` and a
``client_id`` onto a server whose mode promises the gateway stores nothing. A DCR bridge is
the exception by design: its gateway-hosted sign-in registers one shared client bound to the
gateway ``/callback`` (the same app an admin may enter by hand), and that client must survive
for every later caller's authorize and token exchange. ``oauth2_flow`` is still never stamped
on a client-forwarded row.
``redirect_uris`` records what the client is bound to so a later origin change can be
detected as a positive mismatch and trigger re-registration instead of stranding the
@ -1842,7 +1888,7 @@ async def _persist_dcr_client_registration(
``update_mcp_server`` merges credential blobs: a re-registered public client must not
inherit the previous client's secret or auth method.
"""
if mcp_server.is_client_forwarded_token:
if mcp_server.is_client_forwarded_token and not mcp_server.is_dcr_bridge:
return "skipped"
try:
@ -1910,20 +1956,23 @@ async def _persist_dcr_client_registration(
):
await _refresh_persisted_dcr_server(stored)
return "reused"
stamped_flow: Final[Literal["authorization_code"] | None] = (
None if mcp_server.is_client_forwarded_token else "authorization_code"
)
updated_row: Final = await update_mcp_server(
prisma_client=prisma_client,
data=(
UpdateMCPServerRequest(
server_id=mcp_server.server_id,
credentials=credentials,
oauth2_flow="authorization_code",
oauth2_flow=stamped_flow,
token_url=mcp_server.token_url,
)
if mcp_server.token_url
else UpdateMCPServerRequest(
server_id=mcp_server.server_id,
credentials=credentials,
oauth2_flow="authorization_code",
oauth2_flow=stamped_flow,
)
),
touched_by="mcp_oauth_dcr",
@ -2211,7 +2260,7 @@ async def register_client_with_server(
)
return dummy_return
bridge_relay: Final = _dcr_bridge_relays_client_registration(resolved_server)
bridge_relay: Final = _register_relays_bridge_client(resolved_server)
if bridge_relay and not client_redirect_uris:
raise HTTPException(
status_code=400,

5
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -2517,7 +2517,9 @@ if MCP_AVAILABLE:
resolved_redirect_uri: Final = f"{get_request_base_url(request)}/callback" if sealed_code else redirect_uri
caller_client_id: Final = sealed_code.client_id if sealed_code else client_id
caller_client_secret: Final = sealed_code.client_secret if sealed_code else client_secret
resolved_client_id: Final = mcp_server.client_id or caller_client_id or ""
resolved_client_id: Final = (
caller_client_id if sealed_code else (mcp_server.client_id or caller_client_id or "")
)
if not resolved_client_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
@ -2541,6 +2543,7 @@ if MCP_AVAILABLE:
refresh_token=refresh_token,
scope=scope,
client_token_endpoint_auth_method=sealed_code.token_endpoint_auth_method if sealed_code else None,
sealed_client=sealed_code is not None,
)
@router.post(

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

@ -5,11 +5,12 @@ import json
import time
from base64 import urlsafe_b64encode
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException, Request
from fastapi import HTTPException, Request, Response
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy._types import LiteLLM_MCPServerTable
from litellm.types.mcp import MCPAuth
@ -26,6 +27,14 @@ if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
class _BridgeRegistrationPayload(TypedDict, total=False):
client_name: ReadOnly[str]
grant_types: ReadOnly[list[str]]
response_types: ReadOnly[list[str]]
token_endpoint_auth_method: ReadOnly[str]
redirect_uris: ReadOnly[list[str]]
def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None):
credential = {"type": "oauth2", "access_token": access_token}
if refresh_token is not None:
@ -5849,7 +5858,12 @@ def test_oauth_authorization_server_404_for_non_bridge_client_forwarded_server()
assert exc.value.status_code == 404
async def _bridge_register_response(server, request_payload, replace_stale_client=False):
async def _bridge_register_response(
server: "MCPServer",
request_payload: _BridgeRegistrationPayload,
replace_stale_client: bool = False,
database_connected: bool = True,
) -> tuple[Response, MagicMock, AsyncMock]:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
register_client_with_server,
)
@ -5879,6 +5893,10 @@ async def _bridge_register_response(server, request_payload, replace_stale_clien
new_callable=AsyncMock,
return_value=False,
),
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
**({"return_value": MagicMock()} if database_connected else {"side_effect": HTTPException(status_code=500)}),
),
):
response = await register_client_with_server(
request=_bridge_mock_request(),
@ -6033,6 +6051,309 @@ async def test_register_bridge_relay_never_persists():
mock_persist.assert_not_called()
_LAPTOP_LOOPBACK_REDIRECT: Final = "http://127.0.0.1:33418/"
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type_value", ["oauth_delegate", "true_passthrough"])
async def test_register_bridge_with_proxy_base_url_registers_gateway_client_not_laptop_loopback(
monkeypatch: pytest.MonkeyPatch, auth_type_value: Literal["oauth_delegate", "true_passthrough"]
) -> None:
"""A laptop MCP client (VS Code) registers with its own loopback redirect. With PROXY_BASE_URL set
the gateway must not relay that registration upstream: it registers its one shared client bound
to the gateway /callback, keeps it, and hands the caller a client bound to the caller's loopback.
Relaying the loopback is the customer failure on upstreams that restrict redirect domains."""
import json
from litellm.types.mcp import MCPAuth
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
response, mock_async_client, mock_persist = await _bridge_register_response(
_bridge_server(auth_type=MCPAuth(auth_type_value)),
{"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]},
)
posted: Final = mock_async_client.post.call_args.kwargs["json"]
assert posted["redirect_uris"] == ["https://litellm.example.com/callback"], posted
assert posted["client_name"] == "LiteLLM Proxy", posted
mock_persist.assert_awaited_once()
payload: Final = json.loads(response.body.decode("utf-8"))
assert payload["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT], payload
@pytest.mark.asyncio
async def test_register_bridge_without_database_still_relays_client_registration(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A database-free gateway cannot keep the shared client past this process, so even with
PROXY_BASE_URL set it relays the caller's own registration instead of stranding the next authorize."""
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
_, mock_async_client, mock_persist = await _bridge_register_response(
_bridge_server(),
{"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]},
database_connected=False,
)
assert mock_async_client.post.call_args.kwargs["json"]["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT]
mock_persist.assert_not_called()
@pytest.mark.asyncio
async def test_register_bridge_without_proxy_base_url_still_relays_client_registration(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Without PROXY_BASE_URL the gateway cannot keep a shared client, so the relay arm remains the
fallback and the caller's own redirect goes upstream as before."""
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
_, mock_async_client, mock_persist = await _bridge_register_response(
_bridge_server(),
{"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]},
)
assert mock_async_client.post.call_args.kwargs["json"]["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT]
mock_persist.assert_not_called()
def _persisted_shared_gateway_client(
server: "MCPServer", client_id: str = "shared-gateway-client"
) -> LiteLLM_MCPServerTable:
request: Final = _bridge_mock_request()
credentials: Final = {
"client_id": client_id,
"client_secret": None,
"token_endpoint_auth_method": None,
"redirect_uris": [f"{request.base_url.rstrip('/')}/callback"],
"dcr_issuer": server.issuer,
"dcr_server_url": server.url,
}
return LiteLLM_MCPServerTable.model_validate({**server.model_dump(exclude_none=True), "credentials": credentials})
@pytest.mark.asyncio
@pytest.mark.parametrize(
("shared_client_persisted", "expected_redirect_uri", "expected_client_id"),
[
(True, "https://litellm.example.com/callback", "shared-gateway-client"),
(False, _LAPTOP_LOOPBACK_REDIRECT, "dcr-client-123"),
],
)
async def test_authorize_clientless_bridge_loads_shared_client_before_choosing_the_arm(
monkeypatch: pytest.MonkeyPatch,
shared_client_persisted: bool,
expected_redirect_uri: str,
expected_client_id: str,
) -> None:
"""This worker's catalog snapshot has no client_id. When another worker already persisted the shared
client, authorize must load it and send it with the gateway /callback upstream, never the caller's
own client_id or loopback. When no shared client exists yet (a client registered before the gateway
kept one), the caller keeps its own upstream binding through the relay."""
from urllib.parse import parse_qs, urlparse
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding")
server: Final = _bridge_server(url="https://mcp.example.com/mcp")
persisted: Final = _persisted_shared_gateway_client(server) if shared_client_persisted else None
prisma_client: Final = MagicMock()
mock_get_mcp_server: Final = AsyncMock(return_value=persisted)
mock_update_server: Final = AsyncMock()
with (
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
return_value=prisma_client,
),
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,
),
patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper",
return_value="mocked_encrypted_state",
),
):
response: Final = await authorize_with_server(
request=_bridge_mock_request(),
mcp_server=server,
client_id="dcr-client-123",
redirect_uri=_LAPTOP_LOOPBACK_REDIRECT,
state="s",
code_challenge="chal",
code_challenge_method="S256",
)
location: Final = response.headers["location"]
assert location.startswith("https://provider.com/oauth/authorize"), location
query: Final = parse_qs(urlparse(location).query)
assert query["redirect_uri"] == [expected_redirect_uri], query
assert query["client_id"] == [expected_client_id], query
@pytest.mark.asyncio
@pytest.mark.parametrize(
("shared_client_persisted", "expected_redirect_uri", "expected_client_id"),
[
(True, "https://litellm.example.com/callback", "shared-gateway-client"),
(False, _LAPTOP_LOOPBACK_REDIRECT, "dcr-client-123"),
],
)
async def test_token_clientless_bridge_loads_shared_client_before_choosing_the_arm(
monkeypatch: pytest.MonkeyPatch,
shared_client_persisted: bool,
expected_redirect_uri: str,
expected_client_id: str,
) -> None:
"""The token leg mirrors authorize: with the shared client persisted it redeems the code with that
client and the gateway /callback, and with none yet it relays the caller's own client and redirect."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding")
server: Final = _bridge_server(url="https://mcp.example.com/mcp")
persisted: Final = _persisted_shared_gateway_client(server) if shared_client_persisted else None
prisma_client: Final = MagicMock()
mock_get_mcp_server: Final = AsyncMock(return_value=persisted)
mock_update_server: Final = AsyncMock()
with (
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma_client),
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,
),
):
data: Final = await _bridge_token_post_data(server, redirect_uri=_LAPTOP_LOOPBACK_REDIRECT)
assert data["redirect_uri"] == expected_redirect_uri, data
assert data["client_id"] == expected_client_id, data
@pytest.mark.asyncio
@pytest.mark.parametrize(
("dcr_bridge", "expected"),
[(True, "persisted"), (False, "skipped")],
)
async def test_persist_dcr_client_for_client_forwarded_server_only_when_it_is_a_bridge(
dcr_bridge: bool,
expected: Literal["persisted", "skipped"],
) -> None:
"""A DCR bridge keeps the shared gateway client on the server row (without stamping oauth2_flow on
a client-forwarded row), while a client-forwarded server without the bridge still stores nothing."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_persist_dcr_client_registration,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.types.mcp import MCPAuth
server: Final = _bridge_server(auth_type=MCPAuth.oauth_delegate, dcr_bridge=dcr_bridge)
prisma: Final = MagicMock()
prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None)
mock_update: Final = AsyncMock(return_value=MagicMock())
with (
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma),
patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update),
patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()),
patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False),
):
result: Final = await _persist_dcr_client_registration(
server,
{"client_id": "shared-gateway-client", "token_endpoint_auth_method": "none"},
"https://litellm.example.com/callback",
replace_stale_client=False,
)
assert result == expected
if expected == "skipped":
mock_update.assert_not_awaited()
return
written: Final = mock_update.await_args.kwargs["data"]
assert written.oauth2_flow is None, written
assert written.credentials["client_id"] == "shared-gateway-client", written
assert server.client_id == "shared-gateway-client"
@pytest.mark.asyncio
@pytest.mark.parametrize("server_has_shared_client", [False, True])
async def test_token_sealed_ephemeral_client_is_not_replaced_by_persisted_gateway_client(
monkeypatch: pytest.MonkeyPatch,
server_has_shared_client: bool,
) -> None:
"""The token leg must use the sealed ephemeral client even if this worker's catalog has a shared client."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com")
monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding")
server: Final = _bridge_server(
url="https://mcp.example.com/mcp",
client_id="shared-gateway-client" if server_has_shared_client else None,
client_secret="shared-secret" if server_has_shared_client else None,
token_endpoint_auth_method="client_secret_basic" if server_has_shared_client else None,
)
fake_http_response: Final = MagicMock()
fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client: Final = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
prisma_client: Final = MagicMock()
mock_get_mcp_server: Final = AsyncMock(
return_value=_persisted_shared_gateway_client(server, client_id="persisted-gateway-client")
)
mock_update_server: Final = AsyncMock()
with (
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
return_value=prisma_client,
),
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,
),
patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
return_value=fake_http_client,
),
):
await exchange_token_with_server(
request=_bridge_mock_request(),
mcp_server=server,
grant_type="authorization_code",
code="upstream-code",
redirect_uri="https://litellm.example.com/callback",
client_id="ephemeral-client",
client_secret="ephemeral-secret" if server_has_shared_client else None,
code_verifier="verifier",
client_token_endpoint_auth_method="client_secret_post" if server_has_shared_client else None,
sealed_client=True,
)
data: Final = fake_http_client.post.call_args.kwargs["data"]
assert data["client_id"] == "ephemeral-client", data
assert data["redirect_uri"] == "https://litellm.example.com/callback", data
if server_has_shared_client:
headers: Final = fake_http_client.post.call_args.kwargs["headers"]
assert data["client_secret"] == "ephemeral-secret", data
assert "Authorization" not in headers, headers
_BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef"

91
tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -3545,6 +3545,95 @@ class TestTemporaryMCPSessionEndpoints:
assert exchange_mock.await_args.kwargs["redirect_uri"] == "https://litellm.example.com/callback"
assert exchange_mock.await_args.kwargs["client_token_endpoint_auth_method"] == "client_secret_basic"
@pytest.mark.asyncio
async def test_mcp_token_prefers_sealed_client_when_server_has_shared_client(
self: "TestTemporaryMCPSessionEndpoints",
) -> None:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
seal_passthrough_authorization_code,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
mcp_token,
)
request: Final = MagicMock()
request.base_url = "https://litellm.example.com/"
request.headers = {}
server: Final = generate_mock_mcp_server_config_record(server_id="server-1").model_copy(
update={
"auth_type": MCPAuth.true_passthrough,
"client_id": "shared-gateway-client",
"client_secret": "shared-secret",
"token_endpoint_auth_method": "client_secret_basic",
"token_url": "https://provider.example/oauth/token",
}
)
admin_auth: Final = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN)
fake_http_response: Final = MagicMock()
fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"}
fake_http_response.raise_for_status = MagicMock()
fake_http_client: Final = MagicMock()
fake_http_client.post = AsyncMock(return_value=fake_http_response)
prisma_client: Final = MagicMock()
mock_get_mcp_server: Final = AsyncMock(return_value=None)
mock_update_server: Final = AsyncMock()
with (
patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation",
return_value=nullcontext(server),
),
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
return_value=prisma_client,
),
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,
),
patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
return_value=fake_http_client,
),
):
sealed: Final = seal_passthrough_authorization_code(
upstream_code="up-code",
client_id="ephemeral-client",
client_secret="ephemeral-secret",
mcp_server_id="server-1",
token_endpoint_auth_method="client_secret_post",
)
result: Final = await mcp_token(
request=request,
server_id="server-1",
user_api_key_dict=admin_auth,
grant_type="authorization_code",
code=sealed,
redirect_uri="https://example.com/callback",
client_id=None,
client_secret=None,
code_verifier="verifier",
refresh_token=None,
scope=None,
)
data: Final = fake_http_client.post.call_args.kwargs["data"]
headers: Final = fake_http_client.post.call_args.kwargs["headers"]
assert data["client_id"] == "ephemeral-client", data
assert data["client_secret"] == "ephemeral-secret", data
assert data["code"] == "up-code", data
assert data["redirect_uri"] == "https://litellm.example.com/callback", data
assert "Authorization" not in headers, headers
assert result.status_code == 200
assert result.body == b'{"access_token":"tok","token_type":"Bearer"}'
@pytest.mark.asyncio
async def test_mcp_token_refresh_grant_never_opens_sealed_code(self):
"""The minted client is unrecoverable outside the single authorization_code flow by
@ -3841,6 +3930,7 @@ class TestTemporaryMCPSessionEndpoints:
refresh_token=None,
scope=None,
client_token_endpoint_auth_method=None,
sealed_client=False,
)
@pytest.mark.asyncio
@ -3895,6 +3985,7 @@ class TestTemporaryMCPSessionEndpoints:
refresh_token="rt-123",
scope=None,
client_token_endpoint_auth_method=None,
sealed_client=False,
)
@pytest.mark.asyncio