diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3d5b7a32580..13116094614 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index ab9d04a4eb4..09d9809e4fb 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -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) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 5dec580c771..aec92c8d0e4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -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.""" 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 f40904e234d..ef8433fdaf0 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 @@ -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