diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index cf74cbd187e..41d0b78b555 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -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: diff --git a/litellm/proxy/management_endpoints/mcp_connector_import.py b/litellm/proxy/management_endpoints/mcp_connector_import.py index 4a86108515d..8120452393c 100644 --- a/litellm/proxy/management_endpoints/mcp_connector_import.py +++ b/litellm/proxy/management_endpoints/mcp_connector_import.py @@ -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: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 576613c70c4..40cc2e57932 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -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 "" ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py index 408674a0466..9b1a0fb4f98 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_connector_import.py @@ -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( diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index ced0986c831..adab3538b58 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -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()