scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

[MCP Gateway] Allow MCP access groups to be added via the config LIT-312 (#12654)

Browse source 
* allow mcp access groups to be added via the config

* fix code and doc quality

* fix mypy

* create helpers:
This commit is contained in:
Jugal D. Bhatt 2025-07-17 02:27:43 +05:30 • committed by GitHub
parent 6e426c8b7c
commit b3a214bce7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 105 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

65
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -273,34 +273,61 @@ class MCPRequestHandler:
all_servers = direct_mcp_servers + access_group_servers
return list(set(all_servers))
@staticmethod
def _get_config_server_ids_for_access_groups(config_mcp_servers, access_groups: List[str]) -> set:
"""
Helper to get server_ids from config-loaded servers that match any of the given access groups.
"""
server_ids = set()
for server_id, server in config_mcp_servers.items():
if server.access_groups:
if any(group in server.access_groups for group in access_groups):
server_ids.add(server_id)
return server_ids
@staticmethod
async def _get_db_server_ids_for_access_groups(prisma_client, access_groups: List[str]) -> set:
"""
Helper to get server_ids from DB servers that match any of the given access groups.
"""
server_ids = set()
if access_groups and prisma_client is not None:
try:
mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many(
where={
"mcp_access_groups": {
"hasSome": access_groups
}
}
)
for server in mcp_servers:
server_ids.add(server.server_id)
except Exception as e:
verbose_logger.debug(f"Error getting MCP servers from access groups: {e}")
return server_ids
@staticmethod
async def _get_mcp_servers_from_access_groups(
access_groups: List[str]
) -> List[str]:
"""
Resolve MCP access groups to server IDs by querying the MCP server table
Resolve MCP access groups to server IDs by querying BOTH the MCP server table (DB) AND config-loaded servers
"""
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
if not access_groups or prisma_client is None:
return []
# Use the new helper for config-loaded servers
server_ids = MCPRequestHandler._get_config_server_ids_for_access_groups(
global_mcp_server_manager.config_mcp_servers, access_groups
)
try:
# Find all MCP servers that have any of the specified access groups
mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many(
where={
"mcp_access_groups": {
"hasSome": access_groups
}
}
)
# Extract server IDs
server_ids = [server.server_id for server in mcp_servers]
return server_ids
except Exception as e:
verbose_logger.debug(f"Error getting MCP servers from access groups: {e}")
return []
# Use the new helper for DB servers
db_server_ids = await MCPRequestHandler._get_db_server_ids_for_access_groups(
prisma_client, access_groups
)
server_ids.update(db_server_ids)
return list(server_ids)
@staticmethod
async def get_mcp_access_groups(

1
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -132,6 +132,7 @@ class MCPServerManager:
spec_version=server_config.get("spec_version", MCPSpecVersion.mar_2025),
auth_type=server_config.get("auth_type", None),
mcp_info=mcp_info,
access_groups=server_config.get("access_groups", None),
)
self.config_mcp_servers[server_id] = new_server
verbose_logger.debug(

45
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -123,34 +123,31 @@ if MCP_AVAILABLE:
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get all available MCP access groups from the database
Get all available MCP access groups from the database AND config
"""
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
if prisma_client is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={
"error": "Database not connected. Connect a database to your proxy"
},
)
access_groups = set()
try:
# Get all MCP servers and extract their access groups
mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many()
# Extract all unique access groups
access_groups = set()
for server in mcp_servers:
if server.mcp_access_groups:
access_groups.update(server.mcp_access_groups)
# Convert to sorted list
access_groups_list = sorted(list(access_groups))
return {"access_groups": access_groups_list}
except Exception as e:
verbose_proxy_logger.debug(f"Error getting MCP access groups: {e}")
return {"access_groups": []}
# Get from config-loaded servers
for server in global_mcp_server_manager.config_mcp_servers.values():
if server.access_groups:
access_groups.update(server.access_groups)
# Get from DB
if prisma_client is not None:
try:
mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many()
for server in mcp_servers:
if hasattr(server, 'mcp_access_groups') and server.mcp_access_groups:
access_groups.update(server.mcp_access_groups)
except Exception as e:
verbose_proxy_logger.debug(f"Error getting MCP access groups: {e}")
# Convert to sorted list
access_groups_list = sorted(list(access_groups))
return {"access_groups": access_groups_list}
## FastAPI Routes
@router.get(

1
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -27,4 +27,5 @@ class MCPServer(BaseModel):
command: Optional[str] = None
args: Optional[List[str]] = None
env: Optional[Dict[str, str]] = None
access_groups: Optional[List[str]] = None
model_config = ConfigDict(arbitrary_types_allowed=True)

36
tests/mcp_tests/test_mcp_server.py
View file

@ -719,3 +719,39 @@ async def test_get_tools_from_mcp_servers():
pytest.fail(f"Unexpected error in tests: {str(e)}")
def test_mcp_server_manager_access_groups_from_config():
"""
Test that access_groups are loaded from config and can be resolved.
"""
test_manager = MCPServerManager()
test_manager.load_servers_from_config({
"config_server": {
"url": "https://config-mcp-server.com/mcp",
"transport": MCPTransport.http,
"access_groups": ["group-a", "group-b"]
},
"other_server": {
"url": "https://other-mcp-server.com/mcp",
"transport": MCPTransport.http,
"access_groups": ["group-b", "group-c"]
}
})
# Check that access_groups are loaded
config_server = next((s for s in test_manager.config_mcp_servers.values() if s.name == "config_server"), None)
assert config_server is not None
assert set(config_server.access_groups) == {"group-a", "group-b"}
# Check that the lookup logic finds the correct server ids
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
# Patch global_mcp_server_manager for this test
import litellm.proxy._experimental.mcp_server.mcp_server_manager as mcp_server_manager_mod
mcp_server_manager_mod.global_mcp_server_manager = test_manager
# Should find config_server for group-a, both for group-b, other_server for group-c
import asyncio
server_ids_a = asyncio.run(MCPRequestHandler._get_mcp_servers_from_access_groups(["group-a"]))
server_ids_b = asyncio.run(MCPRequestHandler._get_mcp_servers_from_access_groups(["group-b"]))
server_ids_c = asyncio.run(MCPRequestHandler._get_mcp_servers_from_access_groups(["group-c"]))
assert any(config_server.server_id == sid for sid in server_ids_a)
assert set(server_ids_b) == set([s.server_id for s in test_manager.config_mcp_servers.values() if "group-b" in s.access_groups])
assert any(s.name == "other_server" and s.server_id in server_ids_c for s in test_manager.config_mcp_servers.values())