mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix: preserve approved BYOM server visibility
This commit is contained in:
parent
d2aee7e659
commit
94fd2bfa68
4 changed files with 166 additions and 31 deletions
|
|
@ -1246,6 +1246,67 @@ class MCPServerManager:
|
|||
"""Return server IDs that bypass per-key restrictions."""
|
||||
return [server.server_id for server in self.get_registry().values() if server.allow_all_keys is True]
|
||||
|
||||
@staticmethod
|
||||
def get_byom_submitted_servers_cache_key(user_id: str) -> str:
|
||||
return f"byom_submitted_servers:{user_id}"
|
||||
|
||||
async def invalidate_byom_submitted_servers_cache(self, user_id: str | None) -> None:
|
||||
if not user_id:
|
||||
return
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
await user_api_key_cache.async_delete_cache(key=self.get_byom_submitted_servers_cache_key(user_id))
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.warning(f"Failed to invalidate BYOM submitted MCP server cache: {str(e)}")
|
||||
|
||||
async def _get_active_submitted_mcp_server_ids_for_user(
|
||||
self, user_api_key_auth: UserAPIKeyAuth | None
|
||||
) -> list[str]:
|
||||
submitter_user_id = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None
|
||||
if not submitter_user_id:
|
||||
return []
|
||||
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
get_active_submitted_mcp_server_ids_for_user,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.warning(f"Failed to load BYOM submitted MCP server cache dependencies: {str(e)}")
|
||||
return []
|
||||
|
||||
byom_cache_key = self.get_byom_submitted_servers_cache_key(submitter_user_id)
|
||||
submitted_server_ids: list[str] | None = None
|
||||
try:
|
||||
cached_submitted_server_ids = await user_api_key_cache.async_get_cache(key=byom_cache_key)
|
||||
if cached_submitted_server_ids is not None:
|
||||
submitted_server_ids = cast(list[str], cached_submitted_server_ids)
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.warning(f"Failed to read BYOM submitted MCP server cache: {str(e)}")
|
||||
|
||||
if submitted_server_ids is None:
|
||||
if prisma_client is None:
|
||||
submitted_server_ids = []
|
||||
else:
|
||||
try:
|
||||
submitted_server_ids = await get_active_submitted_mcp_server_ids_for_user(
|
||||
prisma_client, submitter_user_id
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.warning(f"Failed to read BYOM submitted MCP servers from database: {str(e)}")
|
||||
submitted_server_ids = []
|
||||
try:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=byom_cache_key,
|
||||
value=submitted_server_ids,
|
||||
ttl=60,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.warning(f"Failed to write BYOM submitted MCP server cache: {str(e)}")
|
||||
|
||||
return [server_id for server_id in submitted_server_ids if self.get_mcp_server_by_id(server_id) is not None]
|
||||
|
||||
async def get_allowed_mcp_servers(self, user_api_key_auth: Optional[UserAPIKeyAuth] = None) -> List[str]:
|
||||
"""
|
||||
Get the allowed MCP Servers for the user.
|
||||
|
|
@ -1258,15 +1319,14 @@ class MCPServerManager:
|
|||
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
|
||||
|
||||
allow_all_server_ids = self.get_allow_all_keys_server_ids()
|
||||
submitted_server_ids = await self._get_active_submitted_mcp_server_ids_for_user(user_api_key_auth)
|
||||
|
||||
try:
|
||||
# The key explicitly opted out of every MCP server. Return zero before
|
||||
# layering on allow_all_keys servers so the opt-out is absolute.
|
||||
key_object_permission = user_api_key_auth.object_permission if user_api_key_auth else None
|
||||
if key_object_permission is not None and (
|
||||
SpecialMCPServerNames.no_mcp_servers.value in (key_object_permission.mcp_servers or [])
|
||||
):
|
||||
return []
|
||||
return submitted_server_ids
|
||||
|
||||
# Check if object_permission.mcp_servers is explicitly set
|
||||
has_explicit_object_permission = False
|
||||
|
|
@ -1325,32 +1385,7 @@ class MCPServerManager:
|
|||
]
|
||||
combined_servers.update(delegate_server_ids)
|
||||
|
||||
# BYOM: approved submissions stay visible to the submitter even when
|
||||
# allow_all_keys=false and no access groups were configured at approval.
|
||||
# ponytail: 60-second TTL cache — upgrade to user-level flag if BYOM adoption grows
|
||||
submitter_user_id = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None
|
||||
if submitter_user_id:
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
get_active_submitted_mcp_server_ids_for_user,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
byom_cache_key = f"byom_submitted_servers:{submitter_user_id}"
|
||||
submitted_server_ids = await user_api_key_cache.async_get_cache(key=byom_cache_key)
|
||||
if submitted_server_ids is None:
|
||||
submitted_server_ids = (
|
||||
await get_active_submitted_mcp_server_ids_for_user(prisma_client, submitter_user_id)
|
||||
if prisma_client is not None
|
||||
else []
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=byom_cache_key,
|
||||
value=submitted_server_ids,
|
||||
ttl=60,
|
||||
)
|
||||
combined_servers.update(
|
||||
server_id for server_id in submitted_server_ids if self.get_mcp_server_by_id(server_id) is not None
|
||||
)
|
||||
combined_servers.update(submitted_server_ids)
|
||||
|
||||
if len(combined_servers) == 0:
|
||||
verbose_logger.debug("No allowed MCP Servers found for user api key auth.")
|
||||
|
|
@ -1358,9 +1393,9 @@ class MCPServerManager:
|
|||
except Exception: # noqa: BLE001
|
||||
verbose_logger.exception(
|
||||
"Failed to get allowed MCP servers; team-level object_permission "
|
||||
"grants may be dropped. Falling back to global servers only."
|
||||
"grants may be dropped. Falling back to global and submitted servers."
|
||||
)
|
||||
return allow_all_server_ids
|
||||
return list(dict.fromkeys(allow_all_server_ids + submitted_server_ids))
|
||||
|
||||
async def resolve_toolset_tool_permissions(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1158,6 +1158,7 @@ if MCP_AVAILABLE:
|
|||
server_id,
|
||||
touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
|
||||
)
|
||||
await global_mcp_server_manager.invalidate_byom_submitted_servers_cache(approved.submitted_by)
|
||||
await global_mcp_server_manager.reload_servers_from_database()
|
||||
|
||||
return _redact_mcp_credentials(approved)
|
||||
|
|
|
|||
|
|
@ -3198,6 +3198,100 @@ class TestMCPServerManager:
|
|||
assert result == []
|
||||
mock_inner.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_mcp_servers_sentinel_keeps_submitted_byom_servers(self):
|
||||
from litellm.proxy import proxy_server as proxy_server_module
|
||||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
MCPRequestHandler,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
|
||||
|
||||
class _Cache:
|
||||
async def async_get_cache(self, key: str):
|
||||
assert key == "byom_submitted_servers:user-123"
|
||||
return ["submitted-server"]
|
||||
|
||||
manager = MCPServerManager()
|
||||
manager.registry = {
|
||||
"submitted-server": MCPServer(
|
||||
server_id="submitted-server",
|
||||
name="submitted",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
}
|
||||
object_permission = LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="perm_no_mcp",
|
||||
mcp_servers=["no-mcp-servers"],
|
||||
mcp_access_groups=[],
|
||||
)
|
||||
user_api_key_auth = UserAPIKeyAuth(
|
||||
api_key="sk-test",
|
||||
user_id="user-123",
|
||||
object_permission=object_permission,
|
||||
object_permission_id="perm_no_mcp",
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "user_api_key_cache", _Cache()),
|
||||
patch.object(proxy_server_module, "prisma_client", None),
|
||||
patch.object(
|
||||
manager, "get_allow_all_keys_server_ids", return_value=["global-server"]
|
||||
),
|
||||
patch.object(
|
||||
MCPRequestHandler,
|
||||
"get_allowed_mcp_servers",
|
||||
new_callable=AsyncMock,
|
||||
return_value=["leaked-server"],
|
||||
) as mock_inner,
|
||||
):
|
||||
result = await manager.get_allowed_mcp_servers(user_api_key_auth)
|
||||
|
||||
assert result == ["submitted-server"]
|
||||
mock_inner.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_fallback_keeps_submitted_byom_servers(self):
|
||||
from litellm.proxy import proxy_server as proxy_server_module
|
||||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
MCPRequestHandler,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
class _Cache:
|
||||
async def async_get_cache(self, key: str):
|
||||
assert key == "byom_submitted_servers:user-123"
|
||||
return ["submitted-server"]
|
||||
|
||||
manager = MCPServerManager()
|
||||
manager.registry = {
|
||||
"submitted-server": MCPServer(
|
||||
server_id="submitted-server",
|
||||
name="submitted",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
}
|
||||
user_api_key_auth = UserAPIKeyAuth(
|
||||
api_key="sk-test",
|
||||
user_id="user-123",
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "user_api_key_cache", _Cache()),
|
||||
patch.object(proxy_server_module, "prisma_client", None),
|
||||
patch.object(
|
||||
manager, "get_allow_all_keys_server_ids", return_value=["global-server"]
|
||||
),
|
||||
patch.object(
|
||||
MCPRequestHandler,
|
||||
"get_allowed_mcp_servers",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=RuntimeError("permission resolver failed"),
|
||||
),
|
||||
):
|
||||
result = await manager.get_allowed_mcp_servers(user_api_key_auth)
|
||||
|
||||
assert set(result) == {"global-server", "submitted-server"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_anonymous_delegate_requires_oauth2(self):
|
||||
"""Anonymous delegated auth listing should only include oauth2 servers."""
|
||||
|
|
|
|||
|
|
@ -3223,8 +3223,10 @@ class TestMCPApprovalWorkflow:
|
|||
pending_server.approval_status = MCPApprovalStatus.pending_review
|
||||
approved_server = generate_mock_mcp_server_db_record()
|
||||
approved_server.approval_status = MCPApprovalStatus.active
|
||||
approved_server.submitted_by = "submitter-user"
|
||||
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.invalidate_byom_submitted_servers_cache = AsyncMock()
|
||||
mock_manager.reload_servers_from_database = AsyncMock()
|
||||
|
||||
with (
|
||||
|
|
@ -3250,6 +3252,9 @@ class TestMCPApprovalWorkflow:
|
|||
)
|
||||
|
||||
mock_manager.reload_servers_from_database.assert_awaited_once()
|
||||
mock_manager.invalidate_byom_submitted_servers_cache.assert_awaited_once_with(
|
||||
"submitter-user"
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue