From fbdff99d41dbe5ab33362c7f6065d6b0b8f99f92 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 9 Mar 2026 20:36:10 -0700 Subject: [PATCH] fix: inline import, add approval workflow tests, rename Submitted MCPs --- .../mcp_management_endpoints.py | 4 +- .../test_mcp_management_endpoints.py | 247 ++++++++++++++++++ .../src/components/mcp_tools/mcp_servers.tsx | 2 +- 3 files changed, 249 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f20bbe04723..cbf32fbb53b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -18,7 +18,7 @@ import importlib import json import os from dataclasses import dataclass -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Any, Dict, Iterable, List, Literal, Optional from fastapi import ( @@ -712,8 +712,6 @@ if MCP_AVAILABLE: Creates the server with approval_status=pending_review. Requires a team-scoped API key. """ - from datetime import datetime, timezone - if not user_api_key_dict.team_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, 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 e81c6264f7b..ecc2e08855d 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 @@ -1512,3 +1512,250 @@ class TestManagementPayloadValidation: assert len(result) == 1 assert result[0]["server_id"] == "server-1" assert result[0]["status"] == "healthy" + + +class TestMCPApprovalWorkflow: + """Tests for BYOM submission: register, list submissions, approve, reject.""" + + @pytest.mark.asyncio + async def test_register_mcp_server_requires_team_key(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + register_mcp_server, + ) + + payload = NewMCPServerRequest( + alias="My Server", + url="https://example.com/mcp", + transport=MCPTransport.sse, + ) + # No team_id → should raise 400 + user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + team_id=None, + ) + with pytest.raises(HTTPException) as exc_info: + await register_mcp_server(payload=payload, user_api_key_dict=user_auth) + assert exc_info.value.status_code == 400 + assert "team" in str(exc_info.value.detail).lower() + + @pytest.mark.asyncio + async def test_register_mcp_server_sets_pending_review(self): + from litellm.proxy._types import MCPApprovalStatus + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + register_mcp_server, + ) + + payload = NewMCPServerRequest( + alias="My Server", + url="https://example.com/mcp", + transport=MCPTransport.sse, + ) + user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-123", + user_id="user-abc", + ) + created_record = generate_mock_mcp_server_db_record( + alias="My Server", + url="https://example.com/mcp", + ) + created_record.approval_status = MCPApprovalStatus.pending_review + created_record.submitted_by = "user-abc" + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + AsyncMock(return_value=created_record), + ) as mock_create, + ): + result = await register_mcp_server( + payload=payload, user_api_key_dict=user_auth + ) + + # Endpoint sets pending_review before calling create_mcp_server + call_payload = mock_create.call_args[0][1] + assert call_payload.approval_status == MCPApprovalStatus.pending_review + assert call_payload.submitted_by == "user-abc" + assert result is not None + + @pytest.mark.asyncio + async def test_get_submissions_non_admin_forbidden(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_mcp_server_submissions, + ) + + non_admin = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.INTERNAL_USER, + ) + with pytest.raises(HTTPException) as exc_info: + await get_mcp_server_submissions(user_api_key_dict=non_admin) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_get_submissions_admin_returns_summary(self): + from litellm.proxy._types import MCPSubmissionsSummary + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_mcp_server_submissions, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + pending = generate_mock_mcp_server_db_record(alias="Pending") + pending.approval_status = "pending_review" + summary = MCPSubmissionsSummary( + total=1, pending_review=1, active=0, rejected=0, items=[pending] + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_submissions", + AsyncMock(return_value=summary), + ), + ): + result = await get_mcp_server_submissions(user_api_key_dict=admin) + + assert result.total == 1 + assert result.pending_review == 1 + + @pytest.mark.asyncio + async def test_approve_non_pending_server_raises_400(self): + from litellm.proxy._types import MCPApprovalStatus + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + approve_mcp_server_submission, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + active_server = generate_mock_mcp_server_db_record() + active_server.approval_status = MCPApprovalStatus.active + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=active_server), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await approve_mcp_server_submission( + server_id="server-1", user_api_key_dict=admin + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_approve_pending_server_loads_into_registry(self): + from litellm.proxy._types import MCPApprovalStatus + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + approve_mcp_server_submission, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + pending_server = generate_mock_mcp_server_db_record() + pending_server.approval_status = MCPApprovalStatus.pending_review + approved_server = generate_mock_mcp_server_db_record() + approved_server.approval_status = MCPApprovalStatus.active + + mock_manager = MagicMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=pending_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.approve_mcp_server", + AsyncMock(return_value=approved_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + result = await approve_mcp_server_submission( + server_id=pending_server.server_id, user_api_key_dict=admin + ) + + mock_manager.reload_servers_from_database.assert_awaited_once() + assert result is not None + + @pytest.mark.asyncio + async def test_reject_already_rejected_raises_400(self): + from litellm.proxy._types import MCPApprovalStatus, RejectMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + reject_mcp_server_submission, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + rejected_server = generate_mock_mcp_server_db_record() + rejected_server.approval_status = MCPApprovalStatus.rejected + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=rejected_server), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await reject_mcp_server_submission( + server_id="server-1", + payload=RejectMCPServerRequest(review_notes="duplicate"), + user_api_key_dict=admin, + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_reject_active_server_allowed(self): + """Admin can deactivate an already-approved server via the reject endpoint.""" + from litellm.proxy._types import MCPApprovalStatus, RejectMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + reject_mcp_server_submission, + ) + + admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + active_server = generate_mock_mcp_server_db_record() + active_server.approval_status = MCPApprovalStatus.active + now_rejected = generate_mock_mcp_server_db_record() + now_rejected.approval_status = MCPApprovalStatus.rejected + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=active_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.reject_mcp_server", + AsyncMock(return_value=now_rejected), + ), + ): + result = await reject_mcp_server_submission( + server_id=active_server.server_id, + payload=RejectMCPServerRequest(review_notes="policy violation"), + user_api_key_dict=admin, + ) + assert result is not None diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index b09fe475e20..0edf8c2a4c6 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -343,7 +343,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) Connect Semantic Filter Network Settings - {isAdminRole(userRole) && Submitted Tools } + {isAdminRole(userRole) && Submitted MCPs }