Create app.py

app.py

@@ -0,0 +1,1124 @@
+import gradio as gr
+import asyncio
+import json
+import logging
+from typing import List, Dict, Any, Optional, Tuple
+from dataclasses import dataclass, field
+import anthropic
+import aiohttp
+from mcp import ClientSession
+from mcp.client.sse import sse_client
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@dataclass
+class MCPSSEServerConfig:
+    name: str
+    url: str
+    headers: Dict[str, str] = field(default_factory=dict)
+    timeout: int = 50
+    sse_read_timeout: int = 50
+
+class FastAgentMCPClient:
+    def __init__(self, config: MCPSSEServerConfig):
+        self.config = config
+        self.session = None
+        self.tools = []
+        self.connected = False
+        self.client_session = None
+    
+    async def connect(self):
+        """Establish connection using fast-agent-mcp"""
+        try:
+            logger.info(f"Connecting to {self.config.name} using fast-agent-mcp")
+            
+            # Use MCP SSE client from fast-agent-mcp
+            self.client_session = await sse_client(
+                self.config.url,
+                headers=self.config.headers,
+                timeout=self.config.timeout
+            )
+            
+            # Initialize the session
+            await self.client_session.initialize()
+            
+            # Get available tools
+            tools_result = await self.client_session.list_tools()
+            self.tools = tools_result.tools if hasattr(tools_result, 'tools') else []
+            
+            self.connected = True
+            logger.info(f"Successfully connected to {self.config.name} with {len(self.tools)} tools")
+            
+        except Exception as e:
+            logger.error(f"Failed to connect to {self.config.name} using fast-agent-mcp: {e}")
+            # Fallback to manual SSE implementation
+            await self._fallback_connect()
+    
+    async def _fallback_connect(self):
+        """Fallback connection method using manual SSE handling"""
+        try:
+            logger.info(f"Attempting fallback connection for {self.config.name}")
+            
+            # Create HTTP session for manual SSE handling
+            connector = aiohttp.TCPConnector(limit=100, limit_per_host=30)
+            self.session = aiohttp.ClientSession(
+                timeout=aiohttp.ClientTimeout(total=self.config.timeout),
+                connector=connector,
+                connector_owner=True
+            )
+            
+            # Test SSE endpoint
+            await self._test_sse_endpoint()
+            
+            # Try to get server capabilities
+            await self._probe_server_capabilities()
+            
+            self.connected = True
+            logger.info(f"Fallback connection successful for {self.config.name}")
+            
+        except Exception as e:
+            logger.warning(f"Fallback connection failed for {self.config.name}: {e}")
+            # Still mark as connected for graceful degradation
+            self.connected = True
+            self.tools = []
+            logger.info(f"Graceful connection established for {self.config.name} (no tools)")
+    
+    async def _test_sse_endpoint(self):
+        """Test SSE endpoint connectivity"""
+        try:
+            async with self.session.get(
+                self.config.url,
+                headers={
+                    "Accept": "text/event-stream",
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "User-Agent": "Klaide-FastAgent-MCP/1.0.0",
+                    **self.config.headers
+                },
+                timeout=aiohttp.ClientTimeout(total=10),
+                allow_redirects=True,
+                max_redirects=10
+            ) as response:
+                
+                # Log redirect chain if any
+                if response.history:
+                    redirect_chain = " -> ".join([str(h.url) for h in response.history])
+                    logger.info(f"SSE endpoint redirected: {redirect_chain} -> {response.url}")
+                
+                if response.status in [200, 204]:
+                    logger.info(f"SSE endpoint accessible for {self.config.name}")
+                elif response.status == 405:
+                    # Handle 405 gracefully - might be expecting different method
+                    logger.warning(f"SSE endpoint returned 405 for {self.config.name} - continuing gracefully")
+                else:
+                    logger.warning(f"SSE endpoint returned {response.status} for {self.config.name}")
+                
+        except Exception as e:
+            logger.warning(f"SSE endpoint test failed for {self.config.name}: {e}")
+    
+    async def _probe_server_capabilities(self):
+        """Probe server for MCP capabilities"""
+        probe_methods = [
+            self._probe_jsonrpc_post,
+            self._probe_jsonrpc_get,
+            self._probe_rest_api,
+            self._probe_websocket
+        ]
+        
+        for method in probe_methods:
+            try:
+                tools = await method()
+                if tools:
+                    self.tools = tools
+                    logger.info(f"Successfully probed {self.config.name} and found {len(tools)} tools")
+                    return
+            except Exception as e:
+                logger.debug(f"Probe method failed for {self.config.name}: {e}")
+                continue
+        
+        # If all probes fail, set empty tools
+        self.tools = []
+        logger.info(f"All probe methods failed for {self.config.name}, setting empty tools")
+    
+    async def _probe_jsonrpc_post(self):
+        """Probe using JSON-RPC POST method"""
+        # Try to initialize with JSON-RPC
+        init_message = {
+            "jsonrpc": "2.0",
+            "id": 1,
+            "method": "initialize",
+            "params": {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {"tools": {}},
+                "clientInfo": {"name": "klaide-fast-agent", "version": "1.0.0"}
+            }
+        }
+        
+        async with self.session.post(
+            self.config.url,
+            json=init_message,
+            headers={
+                "Content-Type": "application/json",
+                "User-Agent": "Klaide-FastAgent-MCP/1.0.0",
+                **self.config.headers
+            },
+            timeout=aiohttp.ClientTimeout(total=15),
+            allow_redirects=True,
+            max_redirects=10
+        ) as response:
+            
+            if response.status != 200:
+                raise Exception(f"Init failed with status {response.status}")
+            
+            result = await response.json()
+            if "error" in result:
+                raise Exception(f"Init error: {result['error']}")
+        
+        # Get tools list
+        tools_message = {
+            "jsonrpc": "2.0",
+            "id": 2,
+            "method": "tools/list",
+            "params": {}
+        }
+        
+        async with self.session.post(
+            self.config.url,
+            json=tools_message,
+            headers={
+                "Content-Type": "application/json",
+                "User-Agent": "Klaide-FastAgent-MCP/1.0.0",
+                **self.config.headers
+            },
+            timeout=aiohttp.ClientTimeout(total=15),
+            allow_redirects=True,
+            max_redirects=10
+        ) as response:
+            
+            if response.status != 200:
+                raise Exception(f"Tools list failed with status {response.status}")
+            
+            tools_response = await response.json()
+            if "result" in tools_response and "tools" in tools_response["result"]:
+                return tools_response["result"]["tools"]
+            
+        return []
+    
+    async def _probe_jsonrpc_get(self):
+        """Probe using JSON-RPC GET method with query parameters"""
+        try:
+            # Try GET with tools list
+            async with self.session.get(
+                f"{self.config.url}?method=tools/list",
+                headers={
+                    "Accept": "application/json",
+                    "User-Agent": "Klaide-FastAgent-MCP/1.0.0",
+                    **self.config.headers
+                },
+                timeout=aiohttp.ClientTimeout(total=10),
+                allow_redirects=True,
+                max_redirects=10
+            ) as response:
+                
+                if response.status == 200:
+                    result = await response.json()
+                    if "tools" in result:
+                        return result["tools"]
+                    elif "result" in result and "tools" in result["result"]:
+                        return result["result"]["tools"]
+                        
+        except Exception:
+            pass
+        
+        return []
+    
+    async def _probe_rest_api(self):
+        """Probe using REST API endpoints"""
+        rest_endpoints = [
+            f"{self.config.url.rstrip('/sse')}/tools",
+            f"{self.config.url.rstrip('/sse')}/api/tools",
+            f"{self.config.url.rstrip('/sse')}/mcp/tools",
+            f"{self.config.url}/tools"
+        ]
+        
+        for endpoint in rest_endpoints:
+            try:
+                async with self.session.get(
+                    endpoint,
+                    headers={
+                        "Accept": "application/json",
+                        "User-Agent": "Klaide-FastAgent-MCP/1.0.0",
+                        **self.config.headers
+                    },
+                    timeout=aiohttp.ClientTimeout(total=10),
+                    allow_redirects=True,
+                    max_redirects=10
+                ) as response:
+                    
+                    if response.status == 200:
+                        result = await response.json()
+                        if isinstance(result, list):
+                            return result
+                        elif "tools" in result:
+                            return result["tools"]
+                            
+            except Exception:
+                continue
+        
+        return []
+    
+    async def _probe_websocket(self):
+        """Probe using WebSocket connection"""
+        try:
+            # Convert HTTP URL to WebSocket URL
+            ws_url = self.config.url.replace('https://', 'wss://').replace('http://', 'ws://')
+            
+            # This would require websockets library
+            # For now, just return empty list
+            logger.debug(f"WebSocket probe not implemented for {ws_url}")
+            return []
+            
+        except Exception:
+            return []
+    
+    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """Call a tool using the best available method"""
+        if self.client_session:
+            # Use fast-agent-mcp client session
+            try:
+                result = await self.client_session.call_tool(tool_name, arguments)
+                return result.content if hasattr(result, 'content') else result
+            except Exception as e:
+                logger.error(f"Fast-agent tool call failed: {e}")
+                # Fall through to manual method
+        
+        # Fallback to manual tool calling
+        if self.session:
+            message = {
+                "jsonrpc": "2.0",
+                "id": int(asyncio.get_event_loop().time() * 1000),
+                "method": "tools/call",
+                "params": {
+                    "name": tool_name,
+                    "arguments": arguments
+                }
+            }
+            
+            try:
+                async with self.session.post(
+                    self.config.url,
+                    json=message,
+                    headers={
+                        "Content-Type": "application/json",
+                        "User-Agent": "Klaide-FastAgent-MCP/1.0.0",
+                        **self.config.headers
+                    },
+                    timeout=aiohttp.ClientTimeout(total=self.config.sse_read_timeout),
+                    allow_redirects=True,
+                    max_redirects=10
+                ) as response:
+                    
+                    if response.status == 200:
+                        result = await response.json()
+                        if "result" in result:
+                            return result["result"].get("content", [])
+                    
+                    raise Exception(f"Tool call failed with status {response.status}")
+                    
+            except Exception as e:
+                logger.error(f"Manual tool call failed: {e}")
+                return [{"type": "text", "text": f"Error: {str(e)}"}]
+        
+        return [{"type": "text", "text": "Error: No connection available"}]
+    
+    async def close(self):
+        """Close all connections"""
+        if self.client_session:
+            try:
+                await self.client_session.close()
+            except Exception:
+                pass
+            
+        if self.session:
+            try:
+                await self.session.close()
+            except Exception:
+                pass
+        
+        self.connected = False
+
+class MCPChatbot:
+    def __init__(self, anthropic_api_key: str, mcp_servers: Dict[str, MCPSSEServerConfig]):
+        self.anthropic_client = anthropic.Anthropic(api_key=anthropic_api_key)
+        self.mcp_servers = mcp_servers
+        self.clients = {}
+        self.available_tools = {}
+        
+    async def initialize_mcp_servers(self):
+        """Initialize connections to all MCP SSE servers using fast-agent-mcp"""
+        for server_name, server_config in self.mcp_servers.items():
+            try:
+                logger.info(f"Connecting to MCP SSE server: {server_name}")
+                
+                client = FastAgentMCPClient(server_config)
+                await client.connect()
+                
+                self.clients[server_name] = client
+                self.available_tools[server_name] = client.tools
+                
+                if client.connected:
+                    logger.info(f"Successfully connected to {server_name} with {len(client.tools)} tools")
+                else:
+                    logger.warning(f"Partial connection to {server_name}")
+                
+            except Exception as e:
+                logger.error(f"Failed to connect to {server_name}: {e}")
+                continue
+    
+    def format_tools_for_claude(self) -> List[Dict]:
+        """Format tools from MCP for Claude API with enhanced context"""
+        claude_tools = []
+        
+        for server_name, tools in self.available_tools.items():
+            for tool in tools:
+                # Enhanced tool description with server context
+                server_context = ""
+                if server_name == "burp_mcp":
+                    server_context = " (Burp Suite - Web Security Testing)"
+                elif server_name == "viper_mcp":
+                    server_context = " (Metasploit - Penetration Testing)"
+                
+                tool_description = tool.get('description', f"Tool from {server_name}")
+                enhanced_description = f"{tool_description}{server_context}"
+                
+                claude_tool = {
+                    "name": f"{server_name}_{tool.get('name', 'unknown')}",
+                    "description": enhanced_description,
+                    "input_schema": tool.get('inputSchema', {
+                        "type": "object",
+                        "properties": {},
+                        "required": []
+                    }),
+                    "server_context": {
+                        "server_name": server_name,
+                        "server_url": self.clients[server_name].config.url if server_name in self.clients else "",
+                        "capabilities": self._get_server_capabilities(server_name)
+                    }
+                }
+                claude_tools.append(claude_tool)
+        
+        return claude_tools
+    
+    def _get_server_capabilities(self, server_name: str) -> List[str]:
+        """Get capabilities for a specific server"""
+        capabilities_map = {
+            "burp_mcp": [
+                "Web application security testing",
+                "Vulnerability scanning", 
+                "HTTP request/response analysis",
+                "Spider/crawling functionality",
+                "Intruder attacks",
+                "Repeater functionality"
+            ],
+            "viper_mcp": [
+                "Penetration testing",
+                "Exploit development",
+                "Payload generation",
+                "Network reconnaissance", 
+                "Post-exploitation",
+                "Metasploit module execution"
+            ]
+        }
+        return capabilities_map.get(server_name, [])
+    
+    async def execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """Execute tool through appropriate MCP server"""
+        parts = tool_name.split('_', 1)
+        if len(parts) < 2:
+            raise ValueError(f"Invalid tool name format: {tool_name}")
+        
+        server_name = parts[0]
+        actual_tool_name = parts[1]
+        
+        if server_name not in self.clients:
+            raise ValueError(f"Server {server_name} not available")
+        
+        client = self.clients[server_name]
+        
+        if not client.connected:
+            raise ValueError(f"Server {server_name} not connected")
+        
+        try:
+            result = await client.call_tool(actual_tool_name, arguments)
+            return result
+        except Exception as e:
+            logger.error(f"Error executing tool {tool_name}: {e}")
+            return [{"type": "text", "text": f"Error: {str(e)}"}]
+    
+    async def chat(self, message: str, history: List[List[str]]) -> tuple:
+        """Main chat function using messages API with tools context"""
+        try:
+            # Get available tools and their context
+            tools = self.format_tools_for_claude()
+            tools_context = await self._get_tools_context()
+            
+            # Bangun system_content dan messages
+            system_content, messages = await self._build_messages_with_context(message, history, tools_context)
+            
+            # Panggil Claude API
+            response = self.anthropic_client.messages.create(
+                model="claude-3-5-sonnet-20241022",
+                max_tokens=4000,
+                messages=messages,
+                tools=tools,
+                temperature=0.1,
+                system=system_content
+            )
+            
+            # Process response
+            assistant_message = ""
+            
+            for content in response.content:
+                if content.type == "text":
+                    assistant_message += content.text
+                elif content.type == "tool_use":
+                    # Execute tool
+                    try:
+                        tool_result = await self.execute_tool(
+                            content.name, 
+                            content.input
+                        )
+                        
+                        # Format tool result
+                        assistant_message += f"\n\n🔧 **Tool Executed**: {content.name}\n"
+                        
+                        if isinstance(tool_result, list):
+                            for item in tool_result:
+                                if isinstance(item, dict) and item.get("type") == "text":
+                                    assistant_message += f"📊 **Result**: {item.get('text', '')}\n"
+                                else:
+                                    assistant_message += f"📊 **Result**: {json.dumps(item, indent=2, ensure_ascii=False)}\n"
+                        else:
+                            assistant_message += f"📊 **Result**: {json.dumps(tool_result, indent=2, ensure_ascii=False)}\n"
+                            
+                    except Exception as e:
+                        assistant_message += f"\n\n❌ **Tool execution failed**: {content.name}\n"
+                        assistant_message += f"🚫 **Error**: {str(e)}\n"
+            
+            # Update history
+            history.append([message, assistant_message])
+            
+            return history, ""
+            
+        except Exception as e:
+            error_msg = f"❌ **Error**: {str(e)}"
+            history.append([message, error_msg])
+            return history, ""
+    
+    async def _build_messages_with_context(self, message: str, history: List[List[str]], tools_context: str) -> Tuple[str, List[Dict]]:
+        """Build messages array with tools context integration"""
+        # Pisahkan system_content
+        system_content = f"""You are Klaide, a Kali Linux AI Desktop assistant that controls cybersecurity tools through MCP servers.
+You help users perform penetration testing, vulnerability assessment, and security analysis.
+
+{tools_context}
+
+Instructions:
+1. Analyze the user's request in the context of available MCP tools
+2. Use the appropriate tools for cybersecurity tasks
+3. Provide helpful guidance and explanations
+4. Be specific about which Kali Linux tools or techniques to use
+5. Always prioritize security and ethical hacking practices
+6. When using tools, explain what you're doing and why
+
+Available tool format: Use the tools provided in the tools list for executing commands."""
+        messages = []
+        # Tambahkan history dan pesan user seperti biasa, TANPA pesan role "system"
+        if history:
+            for user_msg, assistant_msg in history[-5:]:
+                messages.append({"role": "user", "content": user_msg})
+                if assistant_msg:
+                    messages.append({"role": "assistant", "content": assistant_msg})
+        
+        # Add current message
+        messages.append({"role": "user", "content": message})
+        
+        return system_content, messages
+    
+    async def _get_tools_context(self) -> str:
+        """Get context information from all connected MCP servers"""
+        context_parts = []
+        
+        context_parts.append("=== MCP SERVERS CONTEXT ===")
+        
+        for server_name, client in self.clients.items():
+            if not client.connected:
+                continue
+                
+            context_parts.append(f"\n[{server_name.upper()} SERVER]")
+            context_parts.append(f"URL: {client.config.url}")
+            context_parts.append(f"Status: Connected")
+            
+            # Get tools information
+            tools = self.available_tools.get(server_name, [])
+            context_parts.append(f"Available Tools: {len(tools)}")
+            
+            if tools:
+                context_parts.append("Tools List:")
+                for tool in tools:
+                    tool_name = tool.get('name', 'unknown')
+                    tool_desc = tool.get('description', 'No description')
+                    context_parts.append(f"  - {tool_name}: {tool_desc}")
+                    
+                    # Add input schema info
+                    schema = tool.get('inputSchema', {})
+                    if schema.get('properties'):
+                        props = list(schema['properties'].keys())
+                        context_parts.append(f"    Parameters: {', '.join(props)}")
+            
+            # Get server capabilities
+            if server_name == "burp_mcp":
+                context_parts.append("Capabilities:")
+                context_parts.append("  - Web application security testing")
+                context_parts.append("  - Vulnerability scanning")
+                context_parts.append("  - HTTP request/response analysis")
+                context_parts.append("  - Burp Suite integration")
+                
+            elif server_name == "viper_mcp":
+                context_parts.append("Capabilities:")
+                context_parts.append("  - Penetration testing")
+                context_parts.append("  - Exploit development")
+                context_parts.append("  - Metasploit Framework integration")
+                context_parts.append("  - Payload generation")
+        
+        # Add system status
+        connected_count = sum(1 for client in self.clients.values() if client.connected)
+        total_tools = sum(len(tools) for tools in self.available_tools.values())
+        
+        context_parts.append(f"\n[SYSTEM STATUS]")
+        context_parts.append(f"Connected Servers: {connected_count}/{len(self.clients)}")
+        context_parts.append(f"Total Available Tools: {total_tools}")
+        context_parts.append(f"MCP Client: Fast-Agent-MCP")
+        
+        return "\n".join(context_parts)
+    
+    async def _build_completion_prompt(self, message: str, history: List[List[str]], tools_context: str) -> str:
+        """Build completion prompt with conversation history and tools context (deprecated - now using messages)"""
+        # This method is kept for reference but not used anymore
+        # We now use _build_messages_with_context() instead
+        prompt_parts = []
+        
+        # System prompt
+        prompt_parts.append("You are Klaide, a Kali Linux AI Desktop assistant that controls cybersecurity tools through MCP servers.")
+        prompt_parts.append("You help users perform penetration testing, vulnerability assessment, and security analysis.")
+        prompt_parts.append("")
+        
+        # Tools context
+        prompt_parts.append(tools_context)
+        prompt_parts.append("")
+        
+        # Conversation history
+        if history:
+            prompt_parts.append("=== CONVERSATION HISTORY ===")
+            for user_msg, assistant_msg in history[-5:]:  # Last 5 exchanges
+                prompt_parts.append(f"Human: {user_msg}")
+                prompt_parts.append(f"Assistant: {assistant_msg}")
+                prompt_parts.append("")
+        
+        # Current message
+        prompt_parts.append("=== CURRENT REQUEST ===")
+        prompt_parts.append(f"Human: {message}")
+        prompt_parts.append("")
+        
+        # Instructions
+        prompt_parts.append("=== INSTRUCTIONS ===")
+        prompt_parts.append("1. Analyze the user's request in the context of available MCP tools")
+        prompt_parts.append("2. Use the appropriate tools for cybersecurity tasks")
+        prompt_parts.append("3. Provide helpful cybersecurity guidance and explanations")
+        prompt_parts.append("4. Be specific about which Kali Linux tools or techniques to use")
+        prompt_parts.append("5. Always prioritize security and ethical hacking practices")
+        prompt_parts.append("")
+        prompt_parts.append("Assistant: ")
+        
+        return "\n".join(prompt_parts)
+    
+    async def _extract_tool_calls_from_completion(self, completion_text: str) -> List[Dict[str, Any]]:
+        """Extract tool calls from completion text (deprecated - now using native tool calling)"""
+        # This method is kept for reference but not used anymore
+        # We now use Claude's native tool calling in messages API
+        import re
+        
+        tool_calls = []
+        
+        # Pattern to match TOOL_CALL[tool_name](arguments)
+        pattern = r'TOOL_CALL\[([^\]]+)\]\(([^)]*)\)'
+        matches = re.findall(pattern, completion_text)
+        
+        for match in matches:
+            tool_name = match[0].strip()
+            args_str = match[1].strip()
+            
+            # Parse arguments (simple JSON-like parsing)
+            try:
+                if args_str:
+                    # Handle simple key=value pairs
+                    arguments = {}
+                    if '=' in args_str:
+                        pairs = args_str.split(',')
+                        for pair in pairs:
+                            if '=' in pair:
+                                key, value = pair.split('=', 1)
+                                key = key.strip().strip('"\'')
+                                value = value.strip().strip('"\'')
+                                arguments[key] = value
+                    else:
+                        # Try to parse as JSON
+                        arguments = json.loads(args_str) if args_str else {}
+                else:
+                    arguments = {}
+                
+                tool_calls.append({
+                    'name': tool_name,
+                    'arguments': arguments
+                })
+                
+            except Exception as e:
+                logger.warning(f"Failed to parse tool arguments: {args_str}, error: {e}")
+                # Still add the tool call with empty arguments
+                tool_calls.append({
+                    'name': tool_name,
+                    'arguments': {}
+                })
+        
+        return tool_calls
+    
+    async def get_server_status(self) -> str:
+        """Get comprehensive status of all MCP servers with enhanced context"""
+        status = "📡 **Kali Linux MCP Servers Status (Fast-Agent-MCP + Enhanced Context):**\n\n"
+        
+        if not self.clients:
+            status += "❌ **No servers configured**\n"
+            status += "Please configure at least one MCP server URL in Settings.\n"
+            return status
+        
+        for server_name, client in self.clients.items():
+            status += f"## 🖥️ **{server_name.upper()}**\n"
+            status += f"**URL**: `{client.config.url}`\n"
+            
+            try:
+                if not client.connected:
+                    status += "**Status**: ❌ **DISCONNECTED**\n"
+                    status += "**Connection Method**: Fast-Agent-MCP + Fallback\n"
+                    status += "**Tools**: ⚠️ No tools available\n"
+                    status += "**Context**: Not available for enhanced context\n\n"
+                    continue
+                
+                tools = self.available_tools.get(server_name, [])
+                
+                if not tools:
+                    status += "**Status**: ⚠️ **CONNECTED BUT NO TOOLS**\n"
+                    status += "**Connection Method**: Fast-Agent-MCP (Graceful mode)\n"
+                    status += "**Tools**: 📭 Empty tools list\n"
+                    status += "**Context**: Limited context for AI\n"
+                else:
+                    status += f"**Status**: ✅ **CONNECTED & OPERATIONAL**\n"
+                    
+                    # Determine connection method
+                    if client.client_session:
+                        status += "**Connection Method**: 🚀 Fast-Agent-MCP (Native)\n"
+                    else:
+                        status += "**Connection Method**: 🔄 Fast-Agent-MCP (Fallback)\n"
+                    
+                    status += f"**Tools Available**: 🛠️ **{len(tools)} tools**\n"
+                    status += f"**AI Context**: ✅ **Full context available**\n"
+                    
+                    # Display tools with enhanced context integration
+                    status += "**Tools List** (Available with enhanced context):\n"
+                    for i, tool in enumerate(tools, 1):
+                        tool_name = tool.get('name', 'Unknown')
+                        tool_desc = tool.get('description', 'No description available')
+                        
+                        if len(tool_desc) > 60:
+                            tool_desc = tool_desc[:57] + "..."
+                        
+                        status += f"   `{i:2d}.` **{tool_name}** - {tool_desc}\n"
+                    
+                    # Add server-specific capabilities for enhanced context
+                    capabilities = self._get_server_capabilities(server_name)
+                    status += "**Enhanced Context Capabilities**:\n"
+                    for cap in capabilities:
+                        status += f"   • {cap}\n"
+                
+                # Connection health info
+                status += f"**Timeout**: {client.config.timeout}s\n"
+                status += f"**SSE Timeout**: {client.config.sse_read_timeout}s\n"
+                
+            except Exception as e:
+                status += "**Status**: ❌ **ERROR**\n"
+                status += f"**Error Details**: {str(e)}\n"
+                status += "**Tools**: ⚠️ Unable to retrieve tools\n"
+                status += "**Context**: Error in enhanced context\n"
+            
+            status += "\n" + "─" * 50 + "\n\n"
+        
+        # Add enhanced context summary
+        connected_count = sum(1 for client in self.clients.values() if client.connected)
+        total_tools = sum(len(tools) for tools in self.available_tools.values())
+        
+        status += "## 📊 **Enhanced Context Summary**\n"
+        status += f"**Connected Servers**: {connected_count}/{len(self.clients)}\n"
+        status += f"**Total Available Tools**: {total_tools}\n"
+        status += f"**MCP Client**: Fast-Agent-MCP with enhanced context\n"
+        status += f"**AI Model**: Claude-3.5-Sonnet (Messages API with Context)\n"
+        
+        # Context information
+        if total_tools > 0:
+            status += f"**Context Integration**: ✅ **Full MCP context available**\n"
+            status += f"**AI Features**:\n"
+            status += f"   • Rich server context in system messages\n"
+            status += f"   • Native tool calling with Claude Messages API\n"
+            status += f"   • Enhanced cybersecurity guidance\n"
+            status += f"   • Conversation history integration\n"
+            status += f"   • Server capability awareness\n"
+        else:
+            status += f"**Context Integration**: ⚠️ **Limited context**\n"
+        
+        if connected_count == 0:
+            status += "**Overall Status**: ❌ **NO SERVERS OPERATIONAL**\n"
+            status += "**Action Required**: Check server URLs and network connectivity\n"
+        elif connected_count == len(self.clients):
+            status += "**Overall Status**: ✅ **ALL SYSTEMS OPERATIONAL**\n"
+            status += "**Ready**: Klaide enhanced context mode ready for cybersecurity tasks\n"
+        else:
+            status += "**Overall Status**: ⚠️ **PARTIAL CONNECTIVITY**\n"
+            status += "**Action**: Some servers need attention\n"
+        
+        return status
+    
+    async def close_all_connections(self):
+        """Close all MCP connections"""
+        for client in self.clients.values():
+            await client.close()
+
+# Global chatbot instance
+chatbot = None
+
+def create_mcp_servers_config(burp_url: str, viper_url: str) -> Dict[str, MCPSSEServerConfig]:
+    """Create MCP servers configuration from URLs"""
+    servers = {}
+    
+    if burp_url.strip():
+        servers["burp_mcp"] = MCPSSEServerConfig(
+            name="burp_mcp",
+            url=burp_url.strip(),
+            headers={},
+            timeout=50,
+            sse_read_timeout=50
+        )
+    
+    if viper_url.strip():
+        servers["viper_mcp"] = MCPSSEServerConfig(
+            name="viper_mcp", 
+            url=viper_url.strip(),
+            headers={},
+            timeout=50,
+            sse_read_timeout=50
+        )
+    
+    return servers
+
+async def initialize_chatbot(api_key: str, burp_url: str, viper_url: str):
+    """Initialize Klaide with API key and server URLs using fast-agent-mcp"""
+    global chatbot
+    
+    if not api_key:
+        return "❌ Please enter Anthropic API Key"
+    
+    if not burp_url.strip() and not viper_url.strip():
+        return "❌ Please enter at least one MCP server URL"
+    
+    try:
+        mcp_servers = create_mcp_servers_config(burp_url, viper_url)
+        chatbot = MCPChatbot(api_key, mcp_servers)
+        await chatbot.initialize_mcp_servers()
+        
+        connected_servers = [name for name, client in chatbot.clients.items() if client.connected]
+        total_tools = sum(len(tools) for tools in chatbot.available_tools.values())
+        
+        if connected_servers:
+            status_msg = f"✅ Klaide successfully initialized with Fast-Agent-MCP!\n"
+            status_msg += f"🔗 Connected servers: {', '.join(connected_servers)}\n"
+            status_msg += f"🛠️ Total tools available: {total_tools}\n"
+            status_msg += f"🚀 MCP Client: Fast-Agent-MCP with fallback support\n"
+            
+            # Add connection method info
+            for server_name, client in chatbot.clients.items():
+                if client.connected:
+                    method = "Native" if client.client_session else "Fallback"
+                    status_msg += f"📡 {server_name}: {method} connection\n"
+            
+            return status_msg
+        else:
+            return "⚠️ Klaide initialized but no servers connected. Please check your URLs."
+            
+    except Exception as e:
+        return f"❌ Initialization error: {str(e)}"
+
+async def chat_wrapper(message, history):
+    """Wrapper for chat function"""
+    if not chatbot:
+        history.append([message, "❌ Klaide not initialized. Please enter API Key first."])
+        return history, ""
+    
+    return await chatbot.chat(message, history)
+
+async def get_status():
+    """Get server status"""
+    if not chatbot:
+        return "❌ Klaide not initialized"
+    
+    return await chatbot.get_server_status()
+
+async def cleanup():
+    """Cleanup function"""
+    global chatbot
+    if chatbot:
+        await chatbot.close_all_connections()
+
+# Gradio Interface
+def create_interface():
+    with gr.Blocks(title="Klaide (Kali Linux AI Desktop)", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🐉 Klaide (**Kali Linux AI Desktop**)")
+        gr.Markdown("Controlling Kali Linux Desktop with AI using MCP Server.")
+        
+        with gr.Tab("💬 Console"):
+            with gr.Row():
+                with gr.Column(scale=3):
+                    chatbot_ui = gr.Chatbot(
+                        label="Klaide Console",
+                        height=500,
+                        show_copy_button=True,
+                        avatar_images=("assets/user.png", "assets/csalab.png")
+                    )
+                    
+                    with gr.Row():
+                        msg = gr.Textbox(
+                            placeholder="Ask Klaide to control your Kali Linux tools...",
+                            label="Command Prompt",
+                            scale=4
+                        )
+                        send_btn = gr.Button("Send", scale=1, variant="primary")
+                
+                with gr.Column(scale=1):
+                    gr.Markdown("### 🔧 Controls")
+                    clear_btn = gr.Button("Clear Chat", variant="secondary")
+                    
+                    gr.Markdown("### 📊 Server Status")
+                    status_btn = gr.Button("Refresh Status")
+                    status_display = gr.Markdown("Status will be displayed here...")
+        
+        with gr.Tab("⚙️ Settings"):
+            gr.Markdown("## Setup Configuration")
+            
+            with gr.Row():
+                with gr.Column():
+                    api_key_input = gr.Textbox(
+                        label="Anthropic API Key",
+                        type="password",
+                        placeholder="sk-ant-...",
+                        info="Required: Your Anthropic Claude API key"
+                    )
+                    
+                    burp_url_input = gr.Textbox(
+                        label="Burp MCP Server URL",
+                        placeholder="https://burp.csalab.app/sse",
+                        value="https://burp.csalab.app/sse",
+                        info="Optional: URL for Burp Suite MCP server"
+                    )
+                    
+                    viper_url_input = gr.Textbox(
+                        label="Viper MCP Server URL", 
+                        placeholder="https://msf.csalab.app/your-id/sse",
+                        value="https://msf.csalab.app/02b77a05348211f0/sse",
+                        info="Optional: URL for Metasploit MCP server"
+                    )
+                    
+                    with gr.Row():
+                        init_btn = gr.Button("Initialize Klaide", variant="primary", scale=2)
+                        test_urls_btn = gr.Button("Test URLs", variant="secondary", scale=1)
+                    
+                    init_status = gr.Textbox(
+                        label="Klaide Status", 
+                        interactive=False,
+                        lines=4
+                    )
+            
+            gr.Markdown("## Klaide (Kali Linux AI Desktop)")
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                gr.Markdown("### Timeout Configuration")
+                timeout_slider = gr.Slider(
+                    minimum=10,
+                    maximum=120,
+                    value=50,
+                    step=5,
+                    label="Connection Timeout (seconds)",
+                    info="Timeout for server connections and requests"
+                )
+                
+                gr.Markdown("### Custom Headers")
+                custom_headers = gr.Textbox(
+                    label="Custom Headers (JSON format)",
+                    placeholder='{"Authorization": "Bearer token", "X-API-Key": "key"}',
+                    info="Optional: Custom headers for server requests"
+                )
+        
+        # Event handlers
+        def chat_fn(message, history):
+            return asyncio.run(chat_wrapper(message, history))
+        
+        def init_fn(api_key, burp_url, viper_url):
+            return asyncio.run(initialize_chatbot(api_key, burp_url, viper_url))
+        
+        def status_fn():
+            return asyncio.run(get_status())
+        
+        async def test_urls_async(burp_url, viper_url):
+            """Test server URLs connectivity with fast-agent-mcp methods"""
+            results = []
+            
+            # Test Burp URL
+            if burp_url.strip():
+                burp_result = await test_single_url_fast_agent("Burp", burp_url.strip())
+                results.append(burp_result)
+            else:
+                results.append("⏭️ Burp Server: URL not provided")
+            
+            # Test Viper URL  
+            if viper_url.strip():
+                viper_result = await test_single_url_fast_agent("Viper", viper_url.strip())
+                results.append(viper_result)
+            else:
+                results.append("⏭️ Viper Server: URL not provided")
+            
+            return "\n".join(results)
+        
+        async def test_single_url_fast_agent(server_name, url):
+            """Test a single URL using fast-agent-mcp approach"""
+            test_results = []
+            
+            # Method 1: Try fast-agent-mcp native connection
+            try:
+                config = MCPSSEServerConfig(name=f"test_{server_name.lower()}", url=url)
+                test_client = FastAgentMCPClient(config)
+                
+                # Quick connection test (timeout faster for testing)
+                config.timeout = 10
+                await test_client.connect()
+                
+                if test_client.connected:
+                    tool_count = len(test_client.tools)
+                    if test_client.client_session:
+                        test_results.append(f"✅ {server_name} Server: Fast-Agent-MCP native ({tool_count} tools)")
+                    else:
+                        test_results.append(f"✅ {server_name} Server: Fast-Agent-MCP fallback ({tool_count} tools)")
+                    
+                    await test_client.close()
+                    return "\n".join(test_results)
+                else:
+                    test_results.append(f"⚠️ {server_name} Server: Fast-Agent-MCP failed")
+                
+                await test_client.close()
+                
+            except Exception as e:
+                test_results.append(f"❌ {server_name} Server: Fast-Agent-MCP error - {str(e)[:50]}...")
+            
+            # Method 2: Manual SSE test
+            try:
+                connector = aiohttp.TCPConnector(limit=10, limit_per_host=5)
+                async with aiohttp.ClientSession(
+                    timeout=aiohttp.ClientTimeout(total=10),
+                    connector=connector
+                ) as session:
+                    
+                    # Test SSE endpoint
+                    async with session.get(
+                        url,
+                        headers={
+                            "Accept": "text/event-stream",
+                            "User-Agent": "Klaide-FastAgent-Test/1.0.0"
+                        },
+                        allow_redirects=True,
+                        max_redirects=10
+                    ) as response:
+                        
+                        if response.status in [200, 204]:
+                            test_results.append(f"✅ {server_name} Server: SSE endpoint accessible")
+                        elif response.status == 405:
+                            test_results.append(f"⚠️ {server_name} Server: SSE returns 405 (will use graceful handling)")
+                        else:
+                            test_results.append(f"⚠️ {server_name} Server: SSE returns HTTP {response.status}")
+                
+            except Exception as e:
+                test_results.append(f"❌ {server_name} Server: SSE test failed - {str(e)[:50]}...")
+            
+            return "\n".join(test_results)
+        
+        def test_urls_fn(burp_url, viper_url):
+            return asyncio.run(test_urls_async(burp_url, viper_url))
+        
+        # Bind events
+        send_btn.click(
+            chat_fn,
+            inputs=[msg, chatbot_ui],
+            outputs=[chatbot_ui, msg]
+        )
+        
+        msg.submit(
+            chat_fn,
+            inputs=[msg, chatbot_ui],
+            outputs=[chatbot_ui, msg]
+        )
+        
+        clear_btn.click(
+            lambda: ([], ""),
+            outputs=[chatbot_ui, msg]
+        )
+        
+        init_btn.click(
+            init_fn,
+            inputs=[api_key_input, burp_url_input, viper_url_input],
+            outputs=[init_status]
+        )
+        
+        test_urls_btn.click(
+            test_urls_fn,
+            inputs=[burp_url_input, viper_url_input],
+            outputs=[init_status]
+        )
+        
+        status_btn.click(
+            status_fn,
+            outputs=[status_display]
+        )
+        
+        # Cleanup on app close
+        demo.load(None, None, None)
+    
+    return demo
+
+if __name__ == "__main__":
+    try:
+        # Create and launch interface
+        demo = create_interface()
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=True,
+            debug=True
+        )
+    finally:
+        # Cleanup
+        asyncio.run(cleanup())