Update ui_components.py

ui_components.py

@@ -1,32 +1,62 @@
 """
-UI Components for Universal MCP Client - Enhanced with Inference Provider Support
+UI Components for Universal MCP Client - Fixed with optimal MCP guidance
 """
 import gradio as gr
 from gradio import ChatMessage
 from typing import Tuple, List, Dict, Any
+import os
+import logging
+import traceback
+from openai import OpenAI
 
-from config import AppConfig, CUSTOM_CSS
+from config import AppConfig, CUSTOM_CSS, HF_HUB_AVAILABLE
 from chat_handler import ChatHandler
 from server_manager import ServerManager
 from mcp_client import UniversalMCPClient
 
+# Import HuggingFace Hub for login functionality
+if HF_HUB_AVAILABLE:
+    from huggingface_hub import login, logout, whoami
+    from huggingface_hub.utils import HfHubHTTPError
+
+logger = logging.getLogger(__name__)
+
 class UIComponents:
-    """Manages Gradio UI components and event handlers with Inference Provider support"""
+    """Manages Gradio UI components with improved MCP server management"""
     
     def __init__(self, mcp_client: UniversalMCPClient):
         self.mcp_client = mcp_client
         self.chat_handler = ChatHandler(mcp_client)
         self.server_manager = ServerManager(mcp_client)
+        self.current_user = None
+
+    def _initialize_default_servers(self):
+        """Initialize default MCP servers on app startup"""
+        default_servers = [
+            ("background removal", "ysharma/background-removal-mcp"),
+            ("text to video", "ysharma/ltx-video-distilled"),
+            ("text to speech", "ysharma/Kokoro-TTS-mcp-test"),
+            ("text to image", "ysharma/dalle-3-xl-lora-v2")
+        ]
+        
+        logger.info("🚀 Initializing default MCP servers...")
         
-        # State for current LLM backend
-        self.current_backend = "anthropic"  # "anthropic" or "hf_inference"
-        self.current_provider = None
-        self.current_model = None
+        for server_name, space_id in default_servers:
+            try:
+                status_msg, _ = self.server_manager.add_custom_server(server_name, space_id)
+                if "✅" in status_msg:
+                    logger.info(f"✅ Added default server: {server_name}")
+                else:
+                    logger.warning(f"⚠️ Failed to add default server {server_name}: {status_msg}")
+            except Exception as e:
+                logger.error(f"❌ Error adding default server {server_name}: {e}")
+        
+        logger.info(f"📊 Initialized {len(self.mcp_client.servers)} default servers")
     
     def create_interface(self) -> gr.Blocks:
-        """Create the main Gradio interface with provider selection"""
+        """Create the main Gradio interface with improved layout"""
         with gr.Blocks(
-            title="Universal MCP Client", 
+            title="Universal MCP Client - HF Inference Powered", 
             theme=getattr(gr.themes, AppConfig.GRADIO_THEME.title())(),
             fill_height=True,
             css=CUSTOM_CSS
@@ -39,292 +69,246 @@ class UIComponents:
             chatbot = self._create_main_chat_area()
             
             # Set up event handlers
-            self._setup_event_handlers(chatbot)
+            self._setup_event_handlers(chatbot, demo)
             
         return demo
     
     def _create_sidebar(self):
-        """Create the sidebar with LLM provider selection and server management"""
+        """Create the sidebar with login, provider/model selection, and server management"""
         with gr.Sidebar(elem_id="main-sidebar"):
-            gr.Markdown("# Gradio.chat.app")
+            gr.Markdown("# 🤗 chat.gradio.app")
+            
+            # HuggingFace Login Section
+            self._create_login_section()
             
-            # LLM Backend Selection Section
-            self._create_llm_backend_selection()
+            # Provider and Model Selection with defaults
+            self._create_provider_model_selection()
+            
+            # MCP Server Management
+            self._create_server_management_section()
             
             # Collapsible information section
             with gr.Accordion("📚 Guide & Info", open=False):
                 gr.Markdown("""
                 ## 🎯 How To Use
-                - **Choose LLM Backend**: Select between Anthropic Claude or HuggingFace Inference Providers
-                - **Chat with LLM**: Interact with your selected model
-                - **MCP Integration**: Add MCP servers for enhanced capabilities
-                - **Subscribe**: PRO subscribers get higher usage on various services
+                1. **Login**: Login with your HuggingFace account for API access
+                2. **Add MCP Servers**: Connect to various AI tools on 🤗Hub
+                3. **Enable/Disable Servers**: Use checkboxes to control which servers are active
+                4. **Chat**: Interact with GPT-OSS and use connected MCP Servers
                 
-                ## 💭 New UI Features
-                - **Provider Selection**: Choose from multiple inference providers
-                - **Tool Usage**: See tool calls in collapsible accordions with 🔧 icons
-                - **Results**: View tool results in nested thoughts with 📋 icons
-                - **Media Output**: Images, audio, and videos display separately from text
-                - **Real-time Status**: Watch tools execute with pending/done status indicators
+                ## 💭 Features
+                - **GPT-OSS Models**: OpenAI's latest open-source reasoning models (128k context)
+                - **MCP Integration**: Connect to thousands of AI apps on Hub via MCP protocol
+                - **Multi-Provider**: Access via Cerebras, Fireworks, Together AI, and others
+                - **Media Support**: Automatic embedding of media -- images, audio, and video etc
                 """)
-            
-            # Current backend status
-            self.backend_status = gr.Markdown("**Current Backend**: Not configured")
-            
-            # Server management
-            self._create_server_management_section()
     
-    def _create_llm_backend_selection(self):
-        """Create LLM backend selection interface"""
-        gr.Markdown("## 🤖 LLM Backend Selection")
-        
-        # Radio buttons for backend selection
-        self.backend_radio = gr.Radio(
-            choices=[
-                ("Anthropic Claude Sonnet 4", "anthropic"),
-                ("HuggingFace Inference Providers", "hf_inference")
-            ],
-            value="anthropic",
-            label="Choose LLM Backend",
-            info="Select your preferred language model backend"
-        )
-        
-        # Anthropic configuration section
-        with gr.Column(visible=True, elem_classes="provider-selection anthropic-config") as self.anthropic_config:
-            gr.Markdown("### 🔹 Anthropic Claude Configuration")
-            
-            # Check if API key is available
-            if AppConfig.ANTHROPIC_API_KEY:
-                self.anthropic_status = gr.Markdown("✅ **Status**: API key configured")
-            else:
-                self.anthropic_status = gr.Markdown("""
-                ❌ **Status**: API key not found
-                
-                **Setup Instructions**:
-                1. Go to Space Settings → Secrets
-                2. Add `ANTHROPIC_API_KEY` with your Anthropic API key
-                3. Restart the space
-                """)
-        
-        # HuggingFace Inference Provider configuration section
-        with gr.Column(visible=False, elem_classes="provider-selection hf-config") as self.hf_config:
-            gr.Markdown("### 🔸 HuggingFace Inference Provider Configuration")
-            
-            # Check if HF token is available
-            if AppConfig.HF_TOKEN:
-                self.hf_status = gr.Markdown("✅ **Status**: HF token configured")
-            else:
-                self.hf_status = gr.Markdown("""
-                ❌ **Status**: HF token not found
-                
-                **Setup Instructions**:
-                1. Go to Space Settings → Secrets
-                2. Add `HF_TOKEN` with your HuggingFace token
-                3. Restart the space
-                """)
+    def _create_login_section(self):
+        """Create HuggingFace OAuth login section"""
+        with gr.Group(elem_classes="login-section"):
+            gr.Markdown("## 🔑 Authentication", container=True)
+            self.login_button = gr.LoginButton(
+                value="Sign in with Hugging Face",
+                size="sm"
+            )
+            self.login_status = gr.Markdown("⚪ Please sign in to access Inference Providers", container=True)
+    
+    def _create_provider_model_selection(self):
+        """Create provider and model selection dropdowns with defaults"""
+        with gr.Group(elem_classes="provider-model-selection"):
+            gr.Markdown("## 🚀 Inference Configuration", container=True)
             
-            # Provider dropdown
+            # Provider dropdown with default selection
             provider_choices = list(AppConfig.INFERENCE_PROVIDERS.keys())
             self.provider_dropdown = gr.Dropdown(
                 choices=provider_choices,
-                label="🚀 Select Inference Provider",
-                value=provider_choices[0] if provider_choices else None,
+                label="🔧 Inference Provider",
+                value="cerebras",  # Default to Cerebras
                 info="Choose your preferred inference provider"
             )
             
-            # Model dropdown (initially empty, populated based on provider)
+            # Model dropdown (will be populated based on provider)
             self.model_dropdown = gr.Dropdown(
                 choices=[],
-                label="🤖 Select Model",
+                label="🤖 Model",
                 value=None,
-                info="Choose the specific model to use"
+                info="Select GPT OSS model variant"
             )
             
-            # Configure button
-            self.configure_btn = gr.Button(
-                "Configure Inference Provider", 
-                variant="primary",
-                size="sm"
-            )
-            
-            self.provider_config_status = gr.Textbox(
-                label="Configuration Status",
-                interactive=False,
-                visible=False
-            )
+            # Status display
+            self.api_status = gr.Markdown("⚪ Select provider and model to begin", container=True)
     
     def _create_server_management_section(self):
-        """Create the server management section in sidebar"""
-        gr.Markdown("## 🔧 MCP Server Management")
-        
-        # Server status
-        self.server_count_display = gr.Markdown(f"**Connected Servers**: {len(self.mcp_client.servers)}")
-        
-        if self.mcp_client.servers:
-            server_list = "\n".join([f"• **{name}**" for name in self.mcp_client.servers.keys()])
-            self.server_list_display = gr.Markdown(server_list)
-        else:
-            self.server_list_display = gr.Markdown("*No servers connected*\n\nAdd servers below.")
-        
-        with gr.Accordion("⚙️ Manage MCP Servers", open=False):
-            gr.Markdown("**Add MCP Servers**")
+        """Create the server management section with checkboxes and guidance"""
+        with gr.Group():
+            gr.Markdown("## 🔧 MCP Servers", container=True)
+            
+            # ADDED: Optimal server count guidance
+            gr.Markdown("""
+            <div style="background: #f0f8ff; padding: 10px; border-radius: 5px; border-left: 3px solid #4169e1; margin-bottom: 10px;">
+            <strong>💡 Best Practice:</strong> For optimal performance, we recommend keeping 
+            <strong>3-6 MCP servers</strong> enabled at once. Too many servers can:
+            • Increase context usage (reducing available tokens for conversation)
+            • Potentially confuse the model when selecting tools
+            • Slow down response times
+            
+            You can add more servers but selectively enable only the ones you need for your current task.
+            </div>
+            """, container=True)
+            
+            # Server controls
+            with gr.Row():
+                self.add_server_btn = gr.Button("Add MCP Server", variant="primary", size="sm")
+                self.remove_all_btn = gr.Button("Remove All", variant="secondary", size="sm")
+    
+            # Add a save button (initially hidden)
+            self.save_server_btn = gr.Button("Save Server", variant="primary", size="sm", visible=False)
             
-            # Get MCP spaces count for dropdown label
+            # MCP server selection
             from mcp_spaces_finder import _finder
             spaces = _finder.get_mcp_spaces()
-            spaces_count = len(spaces)
-            
-            # Create MCP spaces dropdown
             self.mcp_dropdown = gr.Dropdown(
                 choices=spaces,
-                label=f"🤖 Select from {spaces_count} Available MCP Servers",
+                label=f"**Available MCP Servers ({len(spaces)}**)",
                 value=None,
-                info="Choose a HuggingFace space that provides MCP server functionality",
-                allow_custom_value=True
+                info="Choose from HuggingFace spaces",
+                allow_custom_value=True,
+                visible=False
             )
             
-            # Server title input
             self.server_name = gr.Textbox(
-                label="Give Server a Title", 
-                placeholder="For Example, Text to Image Generator"
+                label="Server Title", 
+                placeholder="e.g., Text to Image Generator",
+                visible=False
             )
             
-            self.add_server_btn = gr.Button("Add Server", variant="primary")
-            self.add_server_output = gr.Textbox(label="Status", interactive=False, container=False)
-            self.add_server_details = gr.HTML(label="Details")
-            
-            # Refresh button for MCP spaces
-            self.refresh_spaces_btn = gr.Button("🔄 Refresh Available Spaces", variant="secondary", size="sm")
+            # Server status and controls
+            self.server_checkboxes = gr.CheckboxGroup(
+                label="Active Servers (Check to enable)",
+                choices=[],
+                value=[],
+                info="✅ Enabled servers can be used | ⬜ Disabled servers are ignored"
+            )
             
-            self.status_btn = gr.Button("Refresh Status", variant="secondary")  
-            self.status_count = gr.Markdown("**Total MCP Servers**: 0")
-            self.status_output = gr.HTML()
+            self.add_server_output = gr.Markdown("", visible=False, container=True)
     
     def _create_main_chat_area(self) -> gr.Chatbot:
-        """Create the main chat area with ChatMessage support"""
+        """Create the main chat area"""
         with gr.Column(elem_classes="main-content"):
-            # Chatbot takes most of the space - configured for ChatMessage
             chatbot = gr.Chatbot(
-                label="Universal MCP-Powered Multimodal Chatbot",
+                label="Universal MCP-Powered AI Assistant",
                 show_label=False,
-                type="messages",  # Essential for ChatMessage support
-                scale=1,  # Expand to fill available space
+                type="messages",
+                scale=1,
                 show_copy_button=True,
                 avatar_images=None,
-                # Initial welcome message with proper ChatMessage format
                 value=[
                     ChatMessage(
                         role="assistant",
-                        content="Welcome! I'm your MCP-powered AI assistant. Please configure your LLM backend in the sidebar to get started. I can help you with various tasks using connected MCP servers!"
+                        content="""Welcome! I'm your MCP-powered AI assistant using OpenAI's GPT-OSS models via HuggingFace Inference Providers.
+                
+                🎉 **Pre-loaded MCP servers ready to use:**
+                - **background removal** - Remove backgrounds from images
+                - **text to video** - Generate videos from text descriptions  
+                - **text to speech** - Convert text to natural speech
+                - **text to image** - Create images from text prompts
+                
+                You can start using these servers right away, add more servers, or remove them as needed. Try asking me to generate an image, create speech, or any other task!"""
                     )
                 ]
             )
-            
-            # Input area at bottom - fixed size
+
+
+
             with gr.Column(scale=0, elem_classes="input-area"):
                 self.chat_input = gr.MultimodalTextbox(
                     interactive=True,
                     file_count="multiple",
-                    placeholder="Enter message or upload files (images, audio, video, documents)...",
+                    placeholder="Enter message or upload files...",
                     show_label=False,
                     sources=["upload", "microphone"],
-                    file_types=None  # Accept all file types
+                    file_types=None
                 )
         
         return chatbot
     
-    def _setup_event_handlers(self, chatbot: gr.Chatbot):
-        """Set up all event handlers for the interface"""
+    def _setup_event_handlers(self, chatbot: gr.Chatbot, demo: gr.Blocks):
+        """Set up all event handlers"""
         
-        # Backend selection event handlers
-        def handle_backend_change(backend_choice):
-            """Handle LLM backend selection change"""
-            self.current_backend = backend_choice
-            
-            # Update visibility of configuration sections
-            anthropic_visible = (backend_choice == "anthropic")
-            hf_visible = (backend_choice == "hf_inference")
-            
-            # Update backend status
-            if backend_choice == "anthropic":
-                if AppConfig.ANTHROPIC_API_KEY:
-                    status = "**Current Backend**: Anthropic Claude Sonnet 4 ✅"
-                else:
-                    status = "**Current Backend**: Anthropic Claude (❌ API key needed)"
+        # OAuth profile handler
+        def handle_oauth_profile(profile: gr.OAuthProfile | None, token: gr.OAuthToken | None):
+            if profile is None:
+                return "⚪ Please sign in to access Inference Providers"
+            
+            logger.info(f"👤 OAuth profile received for user: {profile.name}")
+            
+            if token and token.token:
+                logger.info("🔑 OAuth token received, updating HF client...")
+                os.environ["HF_TOKEN"] = token.token
+                try:
+                    self.mcp_client.hf_client = OpenAI(
+                        base_url="https://router.huggingface.co/v1",
+                        api_key=token.token
+                    )
+                    logger.info("✅ HuggingFace Inference client updated with OAuth token")
+                except Exception as e:
+                    logger.error(f"❌ Failed to update HF client: {e}")
+            
+            return f"✅ Signed in as: **{profile.name}**"
+        
+        # Provider selection with auto-model loading
+        def handle_provider_change(provider_id):
+            if not provider_id:
+                return gr.Dropdown(choices=[], value=None), "⚪ Select provider first"
+            
+            available_models = AppConfig.get_available_models_for_provider(provider_id)
+            model_choices = [(AppConfig.AVAILABLE_MODELS[model]["name"], model) for model in available_models]
+            
+            # Auto-select 120b model if available
+            default_model = "openai/gpt-oss-120b" if "openai/gpt-oss-120b" in available_models else (available_models[0] if available_models else None)
+            
+            # Get context info for status
+            if default_model:
+                model_info = AppConfig.AVAILABLE_MODELS.get(default_model, {})
+                context_length = model_info.get("context_length", 128000)
+                status_msg = f"✅ Provider selected, model auto-selected ({context_length:,} token context)"
             else:
-                if AppConfig.HF_TOKEN:
-                    status = "**Current Backend**: HF Inference Providers (⚙️ Configure provider)"
-                else:
-                    status = "**Current Backend**: HF Inference Providers (❌ HF token needed)"
+                status_msg = "✅ Provider selected, please select a model"
             
             return (
-                gr.update(visible=anthropic_visible),  # anthropic_config
-                gr.update(visible=hf_visible),         # hf_config
-                status                                 # backend_status
+                gr.Dropdown(choices=model_choices, value=default_model, label="🤖 Model"),
+                status_msg
             )
         
-        def handle_provider_change(provider):
-            """Handle inference provider change"""
-            if not provider:
-                return gr.update(choices=[], value=None)
+        # Model selection
+        def handle_model_change(provider_id, model_id):
+            if not provider_id or not model_id:
+                return "⚪ Select both provider and model"
             
-            models = AppConfig.get_provider_models(provider)
-            return gr.update(choices=models, value=models[0] if models else None)
-        
-        def handle_provider_configuration(provider, model):
-            """Handle inference provider configuration"""
-            if not provider or not model:
-                return "❌ Please select both provider and model", gr.update(visible=True)
-            
-            if not AppConfig.HF_TOKEN:
-                return "❌ HF_TOKEN not configured. Please add it to space secrets.", gr.update(visible=True)
+            self.mcp_client.set_model_and_provider(provider_id, model_id)
             
-            # Configure the MCP client with the selected provider
-            success = self.mcp_client.configure_inference_provider(provider, model)
+            # Get model info
+            model_info = AppConfig.AVAILABLE_MODELS.get(model_id, {})
+            context_length = model_info.get("context_length", 128000)
+            active_params = model_info.get("active_params", "N/A")
             
-            if success:
-                self.current_provider = provider
-                self.current_model = model
-                
-                # Update chat handler
-                self.chat_handler.mcp_client = self.mcp_client
-                
-                status_msg = f"✅ Configured {provider} with {model}"
-                
-                # Update backend status
-                backend_status = f"**Current Backend**: {provider}/{model} ✅"
-                
-                return status_msg, gr.update(visible=True)
+            if self.mcp_client.hf_client:
+                return f"✅ Ready! Using {active_params} active params, {context_length:,} token context"
             else:
-                return "❌ Failed to configure inference provider", gr.update(visible=True)
+                return "❌ Please login first"
         
-        # Chat event handlers
+        # Chat handlers
         def submit_message(message, history):
-            """Handle message submission with ChatMessage support"""
-            # Check if backend is properly configured
-            if self.current_backend == "anthropic" and not AppConfig.ANTHROPIC_API_KEY:
-                error_msg = "❌ Please configure Anthropic API key in space settings first."
-                history.append(ChatMessage(role="assistant", content=error_msg))
-                return history, gr.MultimodalTextbox(value=None, interactive=False)
-            
-            if self.current_backend == "hf_inference" and (not self.current_provider or not self.current_model):
-                error_msg = "❌ Please configure HuggingFace inference provider first."
-                history.append(ChatMessage(role="assistant", content=error_msg))
-                return history, gr.MultimodalTextbox(value=None, interactive=False)
-            
             if message and (message.get("text", "").strip() or message.get("files", [])):
-                # Convert existing history to ChatMessage objects if needed
                 converted_history = []
                 for msg in history:
                     if isinstance(msg, dict):
-                        # Convert dict to ChatMessage
                         converted_history.append(ChatMessage(
                             role=msg.get('role', 'assistant'),
                             content=msg.get('content', ''),
                             metadata=msg.get('metadata', None)
                         ))
                     else:
-                        # Already a ChatMessage
                         converted_history.append(msg)
                 
                 new_history, cleared_input = self.chat_handler.process_multimodal_message(message, converted_history)
@@ -332,125 +316,168 @@ class UIComponents:
             return history, gr.MultimodalTextbox(value=None, interactive=False)
         
         def enable_input():
-            """Re-enable input after processing"""
             return gr.MultimodalTextbox(interactive=True)
         
-        # Connect backend selection events
-        self.backend_radio.change(
-            handle_backend_change,
-            inputs=[self.backend_radio],
-            outputs=[self.anthropic_config, self.hf_config, self.backend_status]
+        def show_add_server_fields():
+            return [
+                gr.Dropdown(visible=True),  # mcp_dropdown
+                gr.Textbox(visible=True),   # server_name
+                gr.Button(interactive=False),  # add_server_btn - disable it
+                gr.Button(visible=True)  # save_server_btn - show it
+            ]
+        
+        def hide_add_server_fields():
+            return [
+                gr.Dropdown(visible=False, value=None),  # mcp_dropdown
+                gr.Textbox(visible=False, value=""),   # server_name
+                gr.Button(interactive=True),  # add_server_btn - re-enable it
+                gr.Button(visible=False)  # save_server_btn - hide it
+            ]
+        
+        def handle_add_server(server_title, selected_space):
+            if not server_title or not selected_space:
+                return [
+                    gr.Dropdown(visible=False, value=None),
+                    gr.Textbox(visible=False, value=""),
+                    gr.Button(interactive=True),  # Re-enable add button
+                    gr.Button(visible=False),  # Hide save button
+                    gr.CheckboxGroup(choices=list(self.mcp_client.servers.keys()), 
+                                   value=[name for name, enabled in self.mcp_client.enabled_servers.items() if enabled]),
+                    gr.Markdown("❌ Please provide both server title and space selection", visible=True)
+                ]
+            
+            try:
+                status_msg, _ = self.server_manager.add_custom_server(server_title.strip(), selected_space)
+                
+                # Update checkboxes with all servers
+                server_choices = list(self.mcp_client.servers.keys())
+                enabled_servers = [name for name, enabled in self.mcp_client.enabled_servers.items() if enabled]
+                
+                # Check if we have many servers and show a warning
+                warning_msg = ""
+                if len(enabled_servers) > 6:
+                    warning_msg = "\n\n⚠️ **Note:** You have more than 6 servers enabled. Consider disabling some for better performance."
+                
+                return [
+                    gr.Dropdown(visible=False, value=None),
+                    gr.Textbox(visible=False, value=""),
+                    gr.Button(interactive=True),  # Re-enable add button
+                    gr.Button(visible=False),  # Hide save button
+                    gr.CheckboxGroup(choices=server_choices, value=enabled_servers),
+                    gr.Markdown(status_msg + warning_msg, visible=True)
+                ]
+            
+            except Exception as e:
+                logger.error(f"Error adding server: {e}")
+                return [
+                    gr.Dropdown(visible=False, value=None),
+                    gr.Textbox(visible=False, value=""),
+                    gr.Button(interactive=True),  # Re-enable add button
+                    gr.Button(visible=False),  # Hide save button
+                    gr.CheckboxGroup(choices=list(self.mcp_client.servers.keys()), 
+                                   value=[name for name, enabled in self.mcp_client.enabled_servers.items() if enabled]),
+                    gr.Markdown(f"❌ Error: {str(e)}", visible=True)
+                ]
+        
+        def handle_server_toggle(enabled_servers):
+            """Handle enabling/disabling servers via checkboxes"""
+            # Update enabled status for all servers
+            for server_name in self.mcp_client.servers.keys():
+                self.mcp_client.enable_server(server_name, server_name in enabled_servers)
+            
+            enabled_count = len(enabled_servers)
+            
+            # Provide feedback based on count
+            if enabled_count == 0:
+                message = "ℹ️ No servers enabled - chatbot will use native capabilities only"
+            elif enabled_count <= 6:
+                message = f"✅ {enabled_count} server{'s' if enabled_count != 1 else ''} enabled - optimal configuration"
+            else:
+                message = f"⚠️ {enabled_count} servers enabled - consider reducing to 3-6 for better performance"
+            
+            return gr.Markdown(message, visible=True)
+        
+        def handle_remove_all():
+            """Remove all MCP servers"""
+            count = self.mcp_client.remove_all_servers()
+            return [
+                gr.CheckboxGroup(choices=[], value=[]),
+                gr.Markdown(f"✅ Removed all {count} servers", visible=True)
+            ]
+
+        # Load handler to initialize default mcp servers
+        def initialize_defaults():
+            """Initialize default servers and update UI on app load"""
+            self._initialize_default_servers()
+            
+            # Return updated checkboxes with the default servers
+            server_choices = list(self.mcp_client.servers.keys())
+            enabled_servers = [name for name, enabled in self.mcp_client.enabled_servers.items() if enabled]
+            
+            return gr.CheckboxGroup(
+                choices=server_choices,
+                value=enabled_servers,
+                label=f"Active Servers ({len(server_choices)} loaded)"
+            )
+
+        # Connect OAuth
+        demo.load(
+            fn=handle_oauth_profile,
+            outputs=[self.login_status]
         )
         
-        # Connect provider selection events
+        # Connect provider/model dropdowns with auto-selection on load
+        demo.load(
+            fn=lambda: handle_provider_change("cerebras"),
+            outputs=[self.model_dropdown, self.api_status]
+        )
+
+        # Initialise default mcp server load
+        demo.load(
+            fn=initialize_defaults,
+            outputs=[self.server_checkboxes]
+        )
+
         self.provider_dropdown.change(
             handle_provider_change,
             inputs=[self.provider_dropdown],
-            outputs=[self.model_dropdown]
+            outputs=[self.model_dropdown, self.api_status]
         )
         
-        # Connect provider configuration
-        self.configure_btn.click(
-            handle_provider_configuration,
+        self.model_dropdown.change(
+            handle_model_change,
             inputs=[self.provider_dropdown, self.model_dropdown],
-            outputs=[self.provider_config_status, self.provider_config_status]
-        ).then(
-            lambda: self._update_backend_status(),
-            outputs=[self.backend_status]
+            outputs=[self.api_status]
         )
         
-        # Set up the chat flow - using submit event
-        chat_msg_enter = self.chat_input.submit(
+        # Connect chat
+        chat_submit = self.chat_input.submit(
             submit_message,
             inputs=[self.chat_input, chatbot],
             outputs=[chatbot, self.chat_input]
         )
-        chat_msg_enter.then(enable_input, None, [self.chat_input])
-        
-        # Server management event handlers
-        def update_server_display():
-            """Update the server status display in sidebar"""
-            server_count = len(self.mcp_client.servers)
-            count_text = f"**Connected Servers**: {server_count}"
-            
-            if self.mcp_client.servers:
-                server_list = "\n".join([f"• **{name}**" for name in self.mcp_client.servers.keys()])
-                return count_text, server_list
-            else:
-                return count_text, "*No servers connected*\n\nAdd servers below."
-        
-        def handle_add_server(server_title, selected_space):
-            """Handle adding a server and update displays"""
-            # Check if both title and space are provided
-            if not server_title or not server_title.strip():
-                return "❌ Please provide a server title", "", *update_server_display(), ""
-            
-            if not selected_space:
-                return "❌ Please select a HuggingFace space from the dropdown", "", *update_server_display(), ""
-            
-            # Use the selected space from dropdown
-            status_msg, details_html = self.server_manager.add_custom_server(server_title.strip(), selected_space)
-            
-            # Update sidebar server display
-            count_text, list_text = update_server_display()
-            
-            return status_msg, details_html, count_text, list_text, ""  # Clear server name input
+        chat_submit.then(enable_input, None, [self.chat_input])
         
-        def handle_refresh_status():
-            """Handle refresh status button"""
-            count_text, accordions_html = self.server_manager.get_server_status()
-            return count_text, accordions_html
-        
-        def handle_refresh_spaces():
-            """Handle refresh MCP spaces button"""
-            from mcp_spaces_finder import refresh_mcp_spaces, _finder
-            
-            # Clear cache and get fresh spaces
-            refresh_mcp_spaces()
-            spaces = _finder.get_mcp_spaces()
-            spaces_count = len(spaces)
-            
-            # Update dropdown choices with new count
-            return gr.Dropdown(
-                choices=spaces, 
-                value=None, 
-                label=f"🤖 Select from {spaces_count} Available MCP Server Spaces",
-                info="Choose a HuggingFace space that provides MCP server functionality"
-            )
-        
-        # Connect server management events
+        # Connect server management with proper button state handling
         self.add_server_btn.click(
-            handle_add_server,
-            inputs=[self.server_name, self.mcp_dropdown],
-            outputs=[
-                self.add_server_output, 
-                self.add_server_details, 
-                self.server_count_display, 
-                self.server_list_display, 
-                self.server_name  # Clear server name input
-            ]
+            fn=show_add_server_fields,
+            outputs=[self.mcp_dropdown, self.server_name, self.add_server_btn, self.save_server_btn]
         )
         
-        self.status_btn.click(
-            handle_refresh_status,
-            outputs=[self.status_count, self.status_output]
+        # Connect save button
+        self.save_server_btn.click(
+            fn=handle_add_server,
+            inputs=[self.server_name, self.mcp_dropdown],
+            outputs=[self.mcp_dropdown, self.server_name, self.add_server_btn, self.save_server_btn, self.server_checkboxes, self.add_server_output]
+        )
+                
+        self.server_checkboxes.change(
+            handle_server_toggle,
+            inputs=[self.server_checkboxes],
+            outputs=[self.add_server_output]
         )
         
-        # Connect refresh spaces button
-        self.refresh_spaces_btn.click(
-            handle_refresh_spaces,
-            outputs=[self.mcp_dropdown]
+        self.remove_all_btn.click(
+            handle_remove_all,
+            outputs=[self.server_checkboxes, self.add_server_output]
         )
-    
-    def _update_backend_status(self):
-        """Update backend status based on current configuration"""
-        if self.current_backend == "anthropic":
-            if AppConfig.ANTHROPIC_API_KEY:
-                return "**Current Backend**: Anthropic Claude Sonnet 4 ✅"
-            else:
-                return "**Current Backend**: Anthropic Claude (❌ API key needed)"
-        else:
-            if self.current_provider and self.current_model:
-                return f"**Current Backend**: {self.current_provider}/{self.current_model} ✅"
-            else:
-                return "**Current Backend**: HF Inference Providers (⚙️ Configure provider)"