scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(mcp): forward staged credentials on /mcp-rest/test/connection like /test/tools/list

Browse source 
The connection preview built its temporary MCP client without the credentials the
not-yet-saved server config carries: the Authorization bearer an OAuth2 authorization_code
server had just been granted, the auth_value of an api_key, bearer_token, basic, or
authorization server, and the stored credentials of a saved server being edited. The tools
preview forwarded all three, so the same request succeeded there and failed on the connection
test with the generic "Failed to connect to MCP server" message

Both previews now resolve those credentials through one shared staging step, so they cannot
drift apart again, and the Authorization header is only forwarded upstream when the primary
x-litellm-api-key header carried admission, since otherwise it is the caller's LiteLLM key
This commit is contained in:
mateo-berri 2026-08-29 14:26:48 -07:00
parent 194a3cc202
commit 318b6a4b36
2 changed files with 186 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

85
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -1,11 +1,13 @@
import asyncio
import importlib
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from starlette.datastructures import Headers
from litellm._logging import verbose_logger
from litellm.exceptions import (
@ -1130,6 +1132,45 @@ if MCP_AVAILABLE:
scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None
return client_id, client_secret, scopes
_STAGED_AUTH_VALUE_AUTH_TYPES: Final = frozenset(
(MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization)
)
@dataclass(frozen=True, slots=True)
class _StagedServerTest:
request: NewMCPServerRequest
mcp_auth_header: str | None
oauth2_headers: dict[str, str] | None
def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest:
"""
Resolve the credentials a not-yet-saved server config carries for a preview call.
Both preview endpoints (``/test/connection`` and ``/test/tools/list``) must hand the
temporary client the same credentials, or a server that the saved connection reaches
fine fails one of them.
"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
request: Final = _inherit_credentials_from_existing_server(new_mcp_server_request)
mcp_auth_header: Final = (
request.credentials.get("auth_value")
if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict)
else None
)
# Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY):
# when the primary x-litellm-api-key header is absent, the Authorization value is the
# caller's LiteLLM key, not an upstream token, and must never be forwarded upstream.
oauth2_headers: Final = (
MCPRequestHandler._get_oauth2_headers_from_headers(headers)
if request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
and headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY)
else None
)
return _StagedServerTest(request=request, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers)
async def _execute_with_mcp_client(
request: NewMCPServerRequest,
operation: Callable[..., Awaitable[Mapping[str, object]]],
@ -1339,6 +1380,8 @@ if MCP_AVAILABLE:
},
)
staged: Final = _stage_server_test(new_mcp_server_request, request.headers)
async def _test_connection_operation(client):
async def _noop(session):
return "ok"
@ -1347,8 +1390,10 @@ if MCP_AVAILABLE:
return {"status": "ok"}
return await _execute_with_mcp_client(
new_mcp_server_request,
staged.request,
_test_connection_operation,
mcp_auth_header=staged.mcp_auth_header,
oauth2_headers=staged.oauth2_headers,
raw_headers=_safe_get_request_headers(request),
)
@ -1369,37 +1414,11 @@ if MCP_AVAILABLE:
},
)
new_mcp_server_request = _inherit_credentials_from_existing_server(new_mcp_server_request)
staged: Final = _stage_server_test(new_mcp_server_request, request.headers)
# For OpenAPI spec servers, generate tools from the spec directly
if new_mcp_server_request.spec_path:
return await _preview_openapi_tools(new_mcp_server_request.spec_path)
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
headers: Final = request.headers
mcp_auth_header: str | None = None
if new_mcp_server_request.auth_type in {
MCPAuth.api_key,
MCPAuth.bearer_token,
MCPAuth.basic,
MCPAuth.authorization,
}:
credentials: Final = getattr(new_mcp_server_request, "credentials", None)
if isinstance(credentials, dict):
mcp_auth_header = credentials.get("auth_value")
# Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY):
# when the primary x-litellm-api-key header is absent, the Authorization value is the
# caller's LiteLLM key, not an upstream token, and must never be forwarded upstream.
oauth2_headers: dict[str, str] | None = None
if new_mcp_server_request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and headers.get(
MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY
):
oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
if staged.request.spec_path:
return await _preview_openapi_tools(staged.request.spec_path)
async def _list_tools_operation(client):
async def _list_tools_session_operation(session):
@ -1415,9 +1434,9 @@ if MCP_AVAILABLE:
}
return await _execute_with_mcp_client(
new_mcp_server_request,
staged.request,
_list_tools_operation,
mcp_auth_header=mcp_auth_header,
oauth2_headers=oauth2_headers,
mcp_auth_header=staged.mcp_auth_header,
oauth2_headers=staged.oauth2_headers,
raw_headers=_safe_get_request_headers(request),
)

134
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -464,6 +464,140 @@ class TestTestConnection:
route = _get_route("/mcp-rest/test/connection", "POST")
assert _route_has_dependency(route, user_api_key_auth)
@staticmethod
def _capture_execute(monkeypatch) -> dict:
captured: dict = {}
async def fake_execute(
request,
operation,
mcp_auth_header=None,
oauth2_headers=None,
raw_headers=None,
):
captured["request"] = request
captured["mcp_auth_header"] = mcp_auth_header
captured["oauth2_headers"] = oauth2_headers
return {"status": "ok"}
monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False)
return captured
@staticmethod
def _oauth2_authorization_code_payload(**overrides) -> NewMCPServerRequest:
return NewMCPServerRequest(
server_name="github_mcp",
url="https://api.githubcopilot.com/mcp/",
auth_type=MCPAuth.oauth2,
oauth2_flow="authorization_code",
authorization_url="https://github.com/login/oauth/authorize",
token_url="https://github.com/login/oauth/access_token",
**overrides,
)
@pytest.mark.asyncio
async def test_forwards_staged_oauth2_bearer(self, monkeypatch):
"""The just-authorized upstream token rides the request's Authorization header, exactly
as /test/tools/list receives it; dropping it makes every authorization_code server fail
the connection test that its tools preview passes."""
from litellm.proxy._types import LitellmUserRoles
captured = self._capture_execute(monkeypatch)
request = _build_request(
{"x-litellm-api-key": "sk-admin-session", "authorization": "Bearer upstream-oauth-token"},
path="/mcp-rest/test/connection",
)
result = await rest_endpoints.test_connection(
request,
self._oauth2_authorization_code_payload(),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert result == {"status": "ok"}
assert captured["oauth2_headers"] == {"Authorization": "Bearer upstream-oauth-token"}
assert captured["mcp_auth_header"] is None
@pytest.mark.asyncio
async def test_forwards_staged_auth_value(self, monkeypatch):
from litellm.proxy._types import LitellmUserRoles
captured = self._capture_execute(monkeypatch)
request = _build_request({"x-litellm-api-key": "sk-admin-session"}, path="/mcp-rest/test/connection")
payload = NewMCPServerRequest(
server_name="example",
url="https://example.com/mcp",
auth_type=MCPAuth.bearer_token,
credentials={"auth_value": "upstream-static-token"},
)
result = await rest_endpoints.test_connection(
request,
payload,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert result == {"status": "ok"}
assert captured["mcp_auth_header"] == "upstream-static-token"
assert captured["oauth2_headers"] is None
@pytest.mark.asyncio
async def test_inherits_stored_credentials_of_saved_server(self, monkeypatch):
"""The edit form resends a saved server without its masked credential; the stored one
must be used, as /test/tools/list already does."""
from litellm.proxy._types import LitellmUserRoles
from litellm.types.mcp_server.mcp_server_manager import MCPServer
captured = self._capture_execute(monkeypatch)
saved = MCPServer(
server_id="saved-server-id",
name="example",
url="https://example.com/mcp",
transport="http",
auth_type=MCPAuth.bearer_token,
authentication_token="stored-upstream-token",
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_mcp_server_by_id",
lambda server_id: saved if server_id == "saved-server-id" else None,
)
request = _build_request({"x-litellm-api-key": "sk-admin-session"}, path="/mcp-rest/test/connection")
payload = NewMCPServerRequest(
server_id="saved-server-id",
server_name="example",
url="https://example.com/mcp",
auth_type=MCPAuth.bearer_token,
)
result = await rest_endpoints.test_connection(
request,
payload,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert result == {"status": "ok"}
assert captured["mcp_auth_header"] == "stored-upstream-token"
assert captured["request"].credentials == {"auth_value": "stored-upstream-token"}
@pytest.mark.asyncio
async def test_does_not_forward_authorization_that_satisfied_admission(self, monkeypatch):
"""With no x-litellm-api-key, the Authorization value is the caller's LiteLLM key and
must never reach the upstream."""
from litellm.proxy._types import LitellmUserRoles
captured = self._capture_execute(monkeypatch)
request = _build_request({"authorization": "Bearer sk-litellm-admission-key"}, path="/mcp-rest/test/connection")
result = await rest_endpoints.test_connection(
request,
self._oauth2_authorization_code_payload(),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert result == {"status": "ok"}
assert captured["oauth2_headers"] is None
class TestTestToolsList:
pytestmark = pytest.mark.asyncio