Skip to content

Commit ffd165c

Browse files
Merge pull request #13549 from BerriAI/litellm_mcp_config_fix
[fix] Enhance MCPServerManager with access groups and description support
2 parents a1f6ade + bb32ec9 commit ffd165c

File tree

2 files changed

+90
-0
lines changed

2 files changed

+90
-0
lines changed

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,7 @@ def add_update_server(self, mcp_server: LiteLLM_MCPServerTable):
295295
command=getattr(mcp_server, "command", None),
296296
args=getattr(mcp_server, "args", None) or [],
297297
env=env_dict,
298+
access_groups=getattr(mcp_server, "mcp_access_groups", None),
298299
)
299300
self.registry[mcp_server.server_id] = new_server
300301
verbose_logger.debug(f"Added MCP Server: {name_for_prefix}")
@@ -1050,7 +1051,9 @@ async def get_all_mcp_servers_with_health_and_teams(
10501051
auth_type=_server_config.auth_type,
10511052
created_at=datetime.datetime.now(),
10521053
updated_at=datetime.datetime.now(),
1054+
description=_server_config.mcp_info.get("description") if _server_config.mcp_info else None,
10531055
mcp_info=_server_config.mcp_info,
1056+
mcp_access_groups=_server_config.access_groups or [],
10541057
# Stdio-specific fields
10551058
command=getattr(_server_config, "command", None),
10561059
args=getattr(_server_config, "args", None) or [],

tests/mcp_tests/test_mcp_server.py

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -880,6 +880,93 @@ def test_mcp_server_manager_access_groups_from_config():
880880
assert any(s.name == "other_server" and s.server_id in server_ids_c for s in test_manager.config_mcp_servers.values())
881881

882882

883+
def test_mcp_server_manager_config_integration_with_database():
884+
"""
885+
Test that config-based servers properly integrate with database servers,
886+
specifically testing access_groups and description fields.
887+
"""
888+
import datetime
889+
from litellm.proxy._types import LiteLLM_MCPServerTable
890+
891+
test_manager = MCPServerManager()
892+
893+
# Test 1: Load config with access_groups and description
894+
test_manager.load_servers_from_config({
895+
"config_server_with_groups": {
896+
"url": "https://config-server.com/mcp",
897+
"transport": MCPTransport.http,
898+
"description": "Test config server",
899+
"access_groups": ["fr_staff", "admin"]
900+
}
901+
})
902+
903+
# Verify config server has correct access_groups
904+
config_servers = test_manager.config_mcp_servers
905+
assert len(config_servers) == 1
906+
config_server = next(iter(config_servers.values()))
907+
assert config_server.access_groups == ["fr_staff", "admin"]
908+
assert config_server.mcp_info["description"] == "Test config server"
909+
910+
# Test 2: Create a database server record and test add_update_server method
911+
db_server = LiteLLM_MCPServerTable(
912+
server_id='db-server-123',
913+
server_name='database-server',
914+
url='https://db-server.com/mcp',
915+
transport='http',
916+
spec_version='2025-03-26',
917+
auth_type='none',
918+
description='Database server description',
919+
created_at=datetime.datetime.now(),
920+
updated_at=datetime.datetime.now(),
921+
mcp_access_groups=['db_group', 'test_group']
922+
)
923+
924+
# Test the add_update_server method (this tests our fix)
925+
test_manager.add_update_server(db_server)
926+
927+
# Verify the server was added with correct access_groups
928+
registry = test_manager.get_registry()
929+
assert 'db-server-123' in registry
930+
931+
db_server_in_registry = registry['db-server-123']
932+
assert db_server_in_registry.access_groups == ['db_group', 'test_group']
933+
assert db_server_in_registry.server_name == 'database-server'
934+
935+
# Test 3: Test config server conversion to LiteLLM_MCPServerTable format
936+
# This tests that config servers are properly converted with access_groups and description fields
937+
938+
# Mock user auth to get all servers
939+
from litellm.proxy._types import UserAPIKeyAuth
940+
mock_user_auth = UserAPIKeyAuth(user_role="proxy_admin")
941+
942+
# Mock the get_allowed_mcp_servers to return only config server IDs
943+
# (to avoid database dependency in this test)
944+
async def mock_get_allowed_servers(user_auth=None):
945+
config_server_ids = list(test_manager.config_mcp_servers.keys())
946+
return config_server_ids
947+
948+
test_manager.get_allowed_mcp_servers = mock_get_allowed_servers
949+
950+
# Test the method (this tests our second fix)
951+
import asyncio
952+
servers_list = asyncio.run(test_manager.get_all_mcp_servers_with_health_and_teams(
953+
user_api_key_auth=mock_user_auth
954+
))
955+
956+
# Verify we have the config server properly converted
957+
assert len(servers_list) == 1
958+
959+
# Find the config server in the list
960+
config_server_in_list = servers_list[0]
961+
assert config_server_in_list.server_name == 'config_server_with_groups'
962+
assert config_server_in_list.mcp_access_groups == ["fr_staff", "admin"]
963+
assert config_server_in_list.description == "Test config server"
964+
965+
# Verify the mcp_info is also correct
966+
assert config_server_in_list.mcp_info["description"] == "Test config server"
967+
assert config_server_in_list.mcp_info["server_name"] == "config_server_with_groups"
968+
969+
883970
# Tests for Server Alias Functionality
884971
def test_get_server_prefix_with_alias():
885972
"""

0 commit comments

Comments
 (0)