ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: allow_all_keys to mcp server

Browse Source
This commit is contained in:
Yuta Saito 2026-01-05 15:49:09 +09:00
parent b6664669d2
commit f1d77ae14d
9 changed files with 139 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql Normal file
View File

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "allow_all_keys" BOOLEAN NOT NULL DEFAULT false;

3
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View File

@ -211,6 +211,7 @@ model LiteLLM_MCPServerTable {
authorization_url String?
token_url String?
registration_url String?
allow_all_keys Boolean @default(false)
}
// Generate Tokens for Proxy
@ -748,4 +749,4 @@ model LiteLLM_SkillsTable {
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
}

23
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View File

@ -260,6 +260,7 @@ class MCPServerManager:
allowed_params=server_config.get("allowed_params", None),
access_groups=server_config.get("access_groups", None),
static_headers=server_config.get("static_headers", None),
allow_all_keys=bool(server_config.get("allow_all_keys", False)),
)
self.config_mcp_servers[server_id] = new_server
@ -549,6 +550,7 @@ class MCPServerManager:
access_groups=getattr(mcp_server, "mcp_access_groups", None),
allowed_tools=getattr(mcp_server, "allowed_tools", None),
disallowed_tools=getattr(mcp_server, "disallowed_tools", None),
allow_all_keys=mcp_server.allow_all_keys,
)
return new_server
@ -581,6 +583,14 @@ class MCPServerManager:
all_servers = list(self.get_registry().values())
return {server.server_id for server in all_servers}
def get_allow_all_keys_server_ids(self) -> List[str]:
"""Return server IDs that bypass per-key restrictions."""
return [
server.server_id
for server in self.get_registry().values()
if server.allow_all_keys
]
async def get_allowed_mcp_servers(
self, user_api_key_auth: Optional[UserAPIKeyAuth] = None
) -> List[str]:
@ -593,6 +603,8 @@ class MCPServerManager:
if user_api_key_auth and _user_has_admin_view(user_api_key_auth):
return list(self.get_registry().keys())
allow_all_server_ids = self.get_allow_all_keys_server_ids()
try:
allowed_mcp_servers = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth
@ -600,14 +612,17 @@ class MCPServerManager:
verbose_logger.debug(
f"Allowed MCP Servers for user api key auth: {allowed_mcp_servers}"
)
if len(allowed_mcp_servers) == 0:
combined_servers = set(allowed_mcp_servers)
combined_servers.update(allow_all_server_ids)
if len(combined_servers) == 0:
verbose_logger.debug(
"No allowed MCP Servers found for user api key auth."
)
return allowed_mcp_servers
return list(combined_servers)
except Exception as e:
verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}.")
return []
return allow_all_server_ids
async def get_tools_for_server(self, server_id: str) -> List[MCPTool]:
"""
@ -2238,6 +2253,7 @@ class MCPServerManager:
authorization_url=server.authorization_url,
token_url=server.token_url,
registration_url=server.registration_url,
allow_all_keys=server.allow_all_keys,
)
async def get_all_mcp_servers_with_health_and_teams(
@ -2331,6 +2347,7 @@ class MCPServerManager:
authorization_url=server.authorization_url,
token_url=server.token_url,
registration_url=server.registration_url,
allow_all_keys=server.allow_all_keys,
)
list_mcp_servers.append(mcp_server_table)

5
litellm/proxy/_types.py
View File

@ -1037,6 +1037,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
authorization_url: Optional[str] = None
token_url: Optional[str] = None
registration_url: Optional[str] = None
allow_all_keys: bool = False
@model_validator(mode="before")
@classmethod
@ -1097,6 +1098,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
authorization_url: Optional[str] = None
token_url: Optional[str] = None
registration_url: Optional[str] = None
allow_all_keys: bool = False
@model_validator(mode="before")
@classmethod
@ -1149,6 +1151,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
authorization_url: Optional[str] = None
token_url: Optional[str] = None
registration_url: Optional[str] = None
allow_all_keys: bool = False
class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase):
@ -3821,4 +3824,4 @@ class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase):
class ResponseLiteLLM_ManagedVectorStore(TypedDict, total=False):
vector_store: LiteLLM_ManagedVectorStoresTable
vector_store: LiteLLM_ManagedVectorStoresTable

1
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View File

@ -212,6 +212,7 @@ if MCP_AVAILABLE:
authorization_url=payload.authorization_url,
token_url=payload.token_url,
registration_url=payload.registration_url,
allow_all_keys=payload.allow_all_keys,
)
def get_prisma_client_or_throw(message: str):

3
litellm/proxy/schema.prisma
View File

@ -211,6 +211,7 @@ model LiteLLM_MCPServerTable {
authorization_url String?
token_url String?
registration_url String?
allow_all_keys Boolean @default(false)
}
// Generate Tokens for Proxy
@ -748,4 +749,4 @@ model LiteLLM_SkillsTable {
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
}

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

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

3
schema.prisma
View File

@ -211,6 +211,7 @@ model LiteLLM_MCPServerTable {
authorization_url String?
token_url String?
registration_url String?
allow_all_keys Boolean @default(false)
}
// Generate Tokens for Proxy
@ -748,4 +749,4 @@ model LiteLLM_SkillsTable {
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
}

104
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View File

@ -1057,6 +1057,110 @@ async def test_list_tools_multiple_servers_prefixed_names():
assert names == ["jira-toolA", "zapier-toolA"]
@pytest.mark.asyncio
async def test_mcp_manager_allows_public_servers_without_permissions():
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.proxy._types import MCPTransport
except ImportError:
pytest.skip("MCP server not available")
manager = MCPServerManager()
public_server = MCPServer(
server_id="public",
name="public",
transport=MCPTransport.http,
allow_all_keys=True,
)
manager.registry = {public_server.server_id: public_server}
with patch(
"litellm.proxy.management_endpoints.common_utils._user_has_admin_view",
return_value=False,
), patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPRequestHandler.get_allowed_mcp_servers",
AsyncMock(return_value=[]),
):
allowed = await manager.get_allowed_mcp_servers(UserAPIKeyAuth())
assert allowed == ["public"]
@pytest.mark.asyncio
async def test_mcp_manager_returns_public_when_permission_lookup_fails():
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.proxy._types import MCPTransport
except ImportError:
pytest.skip("MCP server not available")
manager = MCPServerManager()
public_server = MCPServer(
server_id="public",
name="public",
transport=MCPTransport.http,
allow_all_keys=True,
)
manager.registry = {public_server.server_id: public_server}
with patch(
"litellm.proxy.management_endpoints.common_utils._user_has_admin_view",
return_value=False,
), patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPRequestHandler.get_allowed_mcp_servers",
AsyncMock(side_effect=Exception("boom")),
):
allowed = await manager.get_allowed_mcp_servers(UserAPIKeyAuth())
assert allowed == ["public"]
@pytest.mark.asyncio
async def test_mcp_manager_merges_public_and_restricted_servers():
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.proxy._types import MCPTransport
except ImportError:
pytest.skip("MCP server not available")
manager = MCPServerManager()
public_server = MCPServer(
server_id="public",
name="public",
transport=MCPTransport.http,
allow_all_keys=True,
)
scoped_server = MCPServer(
server_id="restricted",
name="restricted",
transport=MCPTransport.http,
)
manager.registry = {
public_server.server_id: public_server,
scoped_server.server_id: scoped_server,
}
with patch(
"litellm.proxy.management_endpoints.common_utils._user_has_admin_view",
return_value=False,
), patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPRequestHandler.get_allowed_mcp_servers",
AsyncMock(return_value=["restricted"]),
):
allowed = await manager.get_allowed_mcp_servers(UserAPIKeyAuth())
assert set(allowed) == {"public", "restricted"}
@pytest.mark.asyncio
async def test_call_mcp_tool_user_unauthorized_access():
"""Test that a user cannot call a tool from a server they don't have access to"""