mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(mcp): harden connector import auth handling and registration
Send a non-Bearer Authorization header verbatim via auth_type authorization instead of wrapping it as a bearer credential, and always drop the Authorization header from static_headers so the plaintext copy cannot shadow the encrypted credential at request time. Register each imported server with the in-memory manager before the best-effort reload, matching the manual add path. Let get_all_mcp_servers propagate read failures instead of returning [], which silently disabled the import dedupe and allowed duplicate imports.
This commit is contained in:
parent
0557a95259
commit
40f5d53c04
5 changed files with 98 additions and 43 deletions
|
|
@ -600,23 +600,19 @@ async def get_all_mcp_servers(
|
|||
NULL approval_status predates the approval workflow, so those rows are kept explicitly rather
|
||||
than dropped by a bare inequality, which SQL evaluates as NULL and would silently hide them.
|
||||
"""
|
||||
try:
|
||||
where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = (
|
||||
{"approval_status": approval_status}
|
||||
if approval_status is not None
|
||||
# mutable-ok: prisma where-inputs must be plain dicts, and both `NOT` and `not` drop
|
||||
# NULL rows (measured), so the OR is the only NULL-preserving way to exclude drafts
|
||||
else {"OR": [{"approval_status": None}, {"approval_status": {"not": MCPApprovalStatus.draft}}]}
|
||||
)
|
||||
mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where)
|
||||
where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = (
|
||||
{"approval_status": approval_status}
|
||||
if approval_status is not None
|
||||
# mutable-ok: prisma where-inputs must be plain dicts, and both `NOT` and `not` drop
|
||||
# NULL rows (measured), so the OR is the only NULL-preserving way to exclude drafts
|
||||
else {"OR": [{"approval_status": None}, {"approval_status": {"not": MCPApprovalStatus.draft}}]}
|
||||
)
|
||||
mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where)
|
||||
|
||||
tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers]
|
||||
for table in tables:
|
||||
decrypt_global_env_var_values(table.env_vars)
|
||||
return tables
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("litellm.proxy._experimental.mcp_server.db.py::get_all_mcp_servers - %s", e)
|
||||
return []
|
||||
tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers]
|
||||
for table in tables:
|
||||
decrypt_global_env_var_values(table.env_vars)
|
||||
return tables
|
||||
|
||||
|
||||
async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None:
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from typing import Final
|
|||
from pydantic import AliasChoices, BaseModel, ConfigDict, Field, ValidationError
|
||||
|
||||
from litellm.proxy._types import MCPApprovalStatus, NewMCPServerRequest
|
||||
from litellm.types.mcp import MCPAuth, MCPCredentials, MCPTransport
|
||||
from litellm.types.mcp import MCPAuth, MCPAuthType, MCPCredentials, MCPTransport
|
||||
|
||||
|
||||
class MCPConnectorEntry(BaseModel):
|
||||
|
|
@ -119,8 +119,7 @@ def _convert_entry(name: str, entry: MCPConnectorEntry) -> ConvertedConnector |
|
|||
return ConnectorConversionError(name=name, error=f"Unsupported connector type '{entry.type}'.")
|
||||
|
||||
transport: Final = MCPTransport.sse if entry_type in _SSE_TYPES else MCPTransport.http
|
||||
token: Final = entry.authorization_token or _header_bearer_token(entry.headers)
|
||||
static_headers: Final = entry.headers if entry.authorization_token else _non_auth_headers(entry.headers)
|
||||
auth: Final = _remote_auth(entry)
|
||||
try:
|
||||
remote_request: Final = NewMCPServerRequest(
|
||||
server_name=sanitized_name,
|
||||
|
|
@ -129,38 +128,37 @@ def _convert_entry(name: str, entry: MCPConnectorEntry) -> ConvertedConnector |
|
|||
approval_status=MCPApprovalStatus.active,
|
||||
transport=transport,
|
||||
url=entry.url,
|
||||
auth_type=MCPAuth.bearer_token if token else MCPAuth.none,
|
||||
credentials=_bearer_credentials(token),
|
||||
static_headers=dict(static_headers) if static_headers is not None else None,
|
||||
auth_type=auth.auth_type,
|
||||
credentials=auth.credentials,
|
||||
static_headers=auth.static_headers,
|
||||
)
|
||||
except ValidationError as e:
|
||||
return ConnectorConversionError(name=name, error=_first_validation_message(e))
|
||||
return ConvertedConnector(name=name, request=remote_request)
|
||||
|
||||
|
||||
def _bearer_credentials(token: str | None) -> MCPCredentials | None:
|
||||
if not token:
|
||||
return None
|
||||
credentials: Final[MCPCredentials] = {"auth_value": token}
|
||||
return credentials
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RemoteAuth:
|
||||
auth_type: MCPAuthType
|
||||
credentials: MCPCredentials | None
|
||||
static_headers: dict[str, str] | None
|
||||
|
||||
|
||||
_AUTHORIZATION_HEADER: Final = "authorization"
|
||||
_BEARER_PREFIX: Final = "bearer "
|
||||
|
||||
|
||||
def _header_bearer_token(headers: Mapping[str, str] | None) -> str | None:
|
||||
values: Final = tuple(value for key, value in (headers or {}).items() if key.lower() == "authorization")
|
||||
if not values:
|
||||
return None
|
||||
value: Final = values[0]
|
||||
return value[len(_BEARER_PREFIX) :] if value.lower().startswith(_BEARER_PREFIX) else value
|
||||
|
||||
|
||||
def _non_auth_headers(headers: Mapping[str, str] | None) -> Mapping[str, str] | None:
|
||||
if headers is None:
|
||||
return None
|
||||
remaining: Final = {key: value for key, value in headers.items() if key.lower() != "authorization"}
|
||||
return remaining or None
|
||||
def _remote_auth(entry: MCPConnectorEntry) -> _RemoteAuth:
|
||||
headers: Final[Mapping[str, str]] = entry.headers or {}
|
||||
header_value: Final = next((value for key, value in headers.items() if key.lower() == _AUTHORIZATION_HEADER), None)
|
||||
remaining: Final = {key: value for key, value in headers.items() if key.lower() != _AUTHORIZATION_HEADER} or None
|
||||
if entry.authorization_token:
|
||||
return _RemoteAuth(MCPAuth.bearer_token, {"auth_value": entry.authorization_token}, remaining)
|
||||
if not header_value:
|
||||
return _RemoteAuth(MCPAuth.none, None, remaining)
|
||||
if header_value.lower().startswith(_BEARER_PREFIX):
|
||||
return _RemoteAuth(MCPAuth.bearer_token, {"auth_value": header_value[len(_BEARER_PREFIX) :]}, remaining)
|
||||
return _RemoteAuth(MCPAuth.authorization, {"auth_value": header_value}, remaining)
|
||||
|
||||
|
||||
def _first_validation_message(error: ValidationError) -> str:
|
||||
|
|
|
|||
|
|
@ -1711,6 +1711,12 @@ if MCP_AVAILABLE:
|
|||
except Exception as e: # noqa: BLE001 # any create failure must become a per-entry error, not a 500
|
||||
verbose_proxy_logger.exception("Error importing mcp server %s: %s", conversion.name, e)
|
||||
return MCPConnectorImportFailure(name=conversion.name, error=str(e))
|
||||
try:
|
||||
await global_mcp_server_manager.add_server(created)
|
||||
except Exception as e: # noqa: BLE001 # the row is committed; the reload after the loop retries registration
|
||||
verbose_proxy_logger.exception(
|
||||
"Imported mcp server %s committed but in-memory registration failed: %s", conversion.name, e
|
||||
)
|
||||
return MCPConnectorImportResult(
|
||||
name=conversion.name, server_id=created.server_id, alias=created.alias or ""
|
||||
)
|
||||
|
|
|
|||
|
|
@ -71,15 +71,24 @@ class TestConvertMcpServersMapping:
|
|||
assert result.request.credentials == {"auth_value": "header-token"}
|
||||
assert result.request.static_headers == {"X-Env": "prod"}
|
||||
|
||||
def test_authorization_header_without_bearer_prefix(self):
|
||||
def test_authorization_header_without_bearer_prefix_is_sent_verbatim(self):
|
||||
result = _single(
|
||||
{"mcpServers": {"srv": {"url": "https://x.example/mcp", "headers": {"authorization": "raw-token"}}}}
|
||||
)
|
||||
assert isinstance(result, ConvertedConnector)
|
||||
assert result.request.auth_type == MCPAuth.bearer_token
|
||||
assert result.request.auth_type == MCPAuth.authorization
|
||||
assert result.request.credentials == {"auth_value": "raw-token"}
|
||||
assert result.request.static_headers is None
|
||||
|
||||
def test_basic_authorization_header_is_sent_verbatim(self):
|
||||
result = _single(
|
||||
{"mcpServers": {"srv": {"url": "https://x.example/mcp", "headers": {"Authorization": "Basic dXNlcjpwdw=="}}}}
|
||||
)
|
||||
assert isinstance(result, ConvertedConnector)
|
||||
assert result.request.auth_type == MCPAuth.authorization
|
||||
assert result.request.credentials == {"auth_value": "Basic dXNlcjpwdw=="}
|
||||
assert result.request.static_headers is None
|
||||
|
||||
def test_authorization_token_wins_over_authorization_header(self):
|
||||
result = _single(
|
||||
{
|
||||
|
|
@ -93,8 +102,9 @@ class TestConvertMcpServersMapping:
|
|||
}
|
||||
)
|
||||
assert isinstance(result, ConvertedConnector)
|
||||
assert result.request.auth_type == MCPAuth.bearer_token
|
||||
assert result.request.credentials == {"auth_value": "explicit-token"}
|
||||
assert result.request.static_headers == {"Authorization": "Bearer header-token"}
|
||||
assert result.request.static_headers is None
|
||||
|
||||
def test_camel_case_authorization_token_alias(self):
|
||||
result = _single(
|
||||
|
|
|
|||
|
|
@ -1960,6 +1960,20 @@ class TestTemporaryMCPSessionEndpoints:
|
|||
where = find_rows.await_args.args[1]
|
||||
assert where == {"OR": [{"approval_status": None}, {"approval_status": {"not": "draft"}}]}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_mcp_servers_propagates_read_failures(self):
|
||||
"""Regression: a swallowed read failure returned [] and silently disabled the bulk-import
|
||||
dedupe, so a flaky DB read turned a re-import into duplicate servers."""
|
||||
from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers
|
||||
|
||||
find_rows = AsyncMock(side_effect=RuntimeError("db down"))
|
||||
with patch( # test-quality-ok: the helper takes its row reader from module scope, matching the suite's pattern
|
||||
"litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows",
|
||||
find_rows,
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="db down"):
|
||||
await get_all_mcp_servers(MagicMock())
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_session_server_id_refuses_an_unknown_caller_supplied_id(self):
|
||||
"""Regression: two concurrent sessions must never land on one id.
|
||||
|
|
@ -6855,6 +6869,7 @@ class TestImportMCPServers:
|
|||
create_mock = AsyncMock(return_value=created)
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.reload_servers_from_database = AsyncMock()
|
||||
mock_manager.add_server = AsyncMock()
|
||||
|
||||
with ExitStack() as stack:
|
||||
for p in self._import_patches([existing], create_mock, mock_manager):
|
||||
|
|
@ -6869,6 +6884,7 @@ class TestImportMCPServers:
|
|||
create_mock.assert_awaited_once()
|
||||
sent_request = create_mock.await_args[0][1]
|
||||
assert sent_request.credentials == {"auth_value": "tok"}
|
||||
mock_manager.add_server.assert_awaited_once_with(created)
|
||||
mock_manager.reload_servers_from_database.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -6891,6 +6907,7 @@ class TestImportMCPServers:
|
|||
create_mock = AsyncMock(return_value=created)
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.reload_servers_from_database = AsyncMock()
|
||||
mock_manager.add_server = AsyncMock()
|
||||
|
||||
with ExitStack() as stack:
|
||||
for p in self._import_patches([], create_mock, mock_manager):
|
||||
|
|
@ -6901,6 +6918,7 @@ class TestImportMCPServers:
|
|||
assert len(result.skipped) == 1
|
||||
assert "Duplicate connector name" in result.skipped[0].reason
|
||||
create_mock.assert_awaited_once()
|
||||
mock_manager.add_server.assert_awaited_once_with(created)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_imports_skips_registry_refresh(self):
|
||||
|
|
@ -6917,6 +6935,7 @@ class TestImportMCPServers:
|
|||
create_mock = AsyncMock()
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.reload_servers_from_database = AsyncMock()
|
||||
mock_manager.add_server = AsyncMock()
|
||||
|
||||
with ExitStack() as stack:
|
||||
for p in self._import_patches([existing], create_mock, mock_manager):
|
||||
|
|
@ -6925,4 +6944,30 @@ class TestImportMCPServers:
|
|||
|
||||
assert result.imported == ()
|
||||
create_mock.assert_not_awaited()
|
||||
mock_manager.add_server.assert_not_awaited()
|
||||
mock_manager.reload_servers_from_database.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_registration_failure_keeps_the_import_result(self):
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
MCPConnectorImportRequest,
|
||||
import_mcp_servers,
|
||||
)
|
||||
|
||||
payload = MCPConnectorImportRequest.model_validate(
|
||||
{"mcpServers": {"new-server": {"url": "https://new.example/mcp"}}}
|
||||
)
|
||||
admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user")
|
||||
created = generate_mock_mcp_server_db_record(server_id="created-1", alias="new_server")
|
||||
create_mock = AsyncMock(return_value=created)
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.reload_servers_from_database = AsyncMock()
|
||||
mock_manager.add_server = AsyncMock(side_effect=RuntimeError("registration boom"))
|
||||
|
||||
with ExitStack() as stack:
|
||||
for p in self._import_patches([], create_mock, mock_manager):
|
||||
stack.enter_context(p)
|
||||
result = await import_mcp_servers(payload=payload, user_api_key_dict=admin)
|
||||
|
||||
assert [entry.name for entry in result.imported] == ["new-server"]
|
||||
mock_manager.reload_servers_from_database.assert_awaited_once()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue