code-puppy 0.0.126__py3-none-any.whl → 0.0.128__py3-none-any.whl

code_puppy/mcp/circuit_breaker.py

@@ -0,0 +1,218 @@
+"""
+Circuit breaker implementation for MCP servers to prevent cascading failures.
+
+This module implements the circuit breaker pattern to protect against cascading
+failures when MCP servers become unhealthy. The circuit breaker has three states:
+- CLOSED: Normal operation, calls pass through
+- OPEN: Calls are blocked and fail fast
+- HALF_OPEN: Limited calls allowed to test recovery
+"""
+
+import asyncio
+import time
+from enum import Enum
+from typing import Any, Callable
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class CircuitState(Enum):
+    """Circuit breaker states."""
+    CLOSED = "closed"      # Normal operation
+    OPEN = "open"          # Blocking calls
+    HALF_OPEN = "half_open"  # Testing recovery
+
+
+class CircuitOpenError(Exception):
+    """Raised when circuit breaker is in OPEN state."""
+    pass
+
+
+class CircuitBreaker:
+    """
+    Circuit breaker to prevent cascading failures in MCP servers.
+    
+    The circuit breaker monitors the success/failure rate of operations and
+    transitions between states to protect the system from unhealthy dependencies.
+    
+    States:
+    - CLOSED: Normal operation, all calls allowed
+    - OPEN: Circuit is open, all calls fail fast with CircuitOpenError
+    - HALF_OPEN: Testing recovery, limited calls allowed
+    
+    State Transitions:
+    - CLOSED → OPEN: After failure_threshold consecutive failures
+    - OPEN → HALF_OPEN: After timeout seconds
+    - HALF_OPEN → CLOSED: After success_threshold consecutive successes
+    - HALF_OPEN → OPEN: After any failure
+    """
+    
+    def __init__(self, failure_threshold: int = 5, success_threshold: int = 2, timeout: int = 60):
+        """
+        Initialize circuit breaker.
+        
+        Args:
+            failure_threshold: Number of consecutive failures before opening circuit
+            success_threshold: Number of consecutive successes needed to close circuit from half-open
+            timeout: Seconds to wait before transitioning from OPEN to HALF_OPEN
+        """
+        self.failure_threshold = failure_threshold
+        self.success_threshold = success_threshold  
+        self.timeout = timeout
+        
+        self._state = CircuitState.CLOSED
+        self._failure_count = 0
+        self._success_count = 0
+        self._last_failure_time = None
+        self._lock = asyncio.Lock()
+        
+        logger.info(
+            f"Circuit breaker initialized: failure_threshold={failure_threshold}, "
+            f"success_threshold={success_threshold}, timeout={timeout}s"
+        )
+    
+    async def call(self, func: Callable, *args, **kwargs) -> Any:
+        """
+        Execute a function through the circuit breaker.
+        
+        Args:
+            func: Function to execute
+            *args: Positional arguments for the function
+            **kwargs: Keyword arguments for the function
+            
+        Returns:
+            Result of the function call
+            
+        Raises:
+            CircuitOpenError: If circuit is in OPEN state
+            Exception: Any exception raised by the wrapped function
+        """
+        async with self._lock:
+            current_state = self._get_current_state()
+            
+            if current_state == CircuitState.OPEN:
+                logger.warning("Circuit breaker is OPEN, failing fast")
+                raise CircuitOpenError("Circuit breaker is open")
+            
+            if current_state == CircuitState.HALF_OPEN:
+                # In half-open state, we're testing recovery
+                logger.info("Circuit breaker is HALF_OPEN, allowing test call")
+        
+        # Execute the function outside the lock to avoid blocking other calls
+        try:
+            result = await func(*args, **kwargs) if asyncio.iscoroutinefunction(func) else func(*args, **kwargs)
+            await self._on_success()
+            return result
+        except Exception as e:
+            await self._on_failure()
+            raise e
+    
+    def record_success(self) -> None:
+        """Record a successful operation."""
+        asyncio.create_task(self._on_success())
+    
+    def record_failure(self) -> None:
+        """Record a failed operation."""
+        asyncio.create_task(self._on_failure())
+    
+    def get_state(self) -> CircuitState:
+        """Get current circuit breaker state."""
+        return self._get_current_state()
+    
+    def is_open(self) -> bool:
+        """Check if circuit breaker is in OPEN state."""
+        return self._get_current_state() == CircuitState.OPEN
+    
+    def is_half_open(self) -> bool:
+        """Check if circuit breaker is in HALF_OPEN state."""
+        return self._get_current_state() == CircuitState.HALF_OPEN
+    
+    def is_closed(self) -> bool:
+        """Check if circuit breaker is in CLOSED state."""
+        return self._get_current_state() == CircuitState.CLOSED
+    
+    def reset(self) -> None:
+        """Reset circuit breaker to CLOSED state and clear counters."""
+        logger.info("Resetting circuit breaker to CLOSED state")
+        self._state = CircuitState.CLOSED
+        self._failure_count = 0
+        self._success_count = 0
+        self._last_failure_time = None
+    
+    def force_open(self) -> None:
+        """Force circuit breaker to OPEN state."""
+        logger.warning("Forcing circuit breaker to OPEN state")
+        self._state = CircuitState.OPEN
+        self._last_failure_time = time.time()
+    
+    def force_close(self) -> None:
+        """Force circuit breaker to CLOSED state and reset counters."""
+        logger.info("Forcing circuit breaker to CLOSED state")
+        self._state = CircuitState.CLOSED
+        self._failure_count = 0
+        self._success_count = 0
+        self._last_failure_time = None
+    
+    def _get_current_state(self) -> CircuitState:
+        """
+        Get the current state, handling automatic transitions.
+        
+        This method handles the automatic transition from OPEN to HALF_OPEN
+        after the timeout period has elapsed.
+        """
+        if self._state == CircuitState.OPEN and self._should_attempt_reset():
+            logger.info("Timeout reached, transitioning from OPEN to HALF_OPEN")
+            self._state = CircuitState.HALF_OPEN
+            self._success_count = 0  # Reset success counter for half-open testing
+        
+        return self._state
+    
+    def _should_attempt_reset(self) -> bool:
+        """Check if enough time has passed to attempt reset from OPEN to HALF_OPEN."""
+        if self._last_failure_time is None:
+            return False
+        
+        return time.time() - self._last_failure_time >= self.timeout
+    
+    async def _on_success(self) -> None:
+        """Handle successful operation."""
+        async with self._lock:
+            current_state = self._get_current_state()
+            
+            if current_state == CircuitState.CLOSED:
+                # Reset failure count on success in closed state
+                if self._failure_count > 0:
+                    logger.debug("Resetting failure count after success")
+                    self._failure_count = 0
+            
+            elif current_state == CircuitState.HALF_OPEN:
+                self._success_count += 1
+                logger.debug(f"Success in HALF_OPEN state: {self._success_count}/{self.success_threshold}")
+                
+                if self._success_count >= self.success_threshold:
+                    logger.info("Success threshold reached, transitioning from HALF_OPEN to CLOSED")
+                    self._state = CircuitState.CLOSED
+                    self._failure_count = 0
+                    self._success_count = 0
+                    self._last_failure_time = None
+    
+    async def _on_failure(self) -> None:
+        """Handle failed operation."""
+        async with self._lock:
+            current_state = self._get_current_state()
+            
+            if current_state == CircuitState.CLOSED:
+                self._failure_count += 1
+                logger.debug(f"Failure in CLOSED state: {self._failure_count}/{self.failure_threshold}")
+                
+                if self._failure_count >= self.failure_threshold:
+                    logger.warning("Failure threshold reached, transitioning from CLOSED to OPEN")
+                    self._state = CircuitState.OPEN
+                    self._last_failure_time = time.time()
+            
+            elif current_state == CircuitState.HALF_OPEN:
+                logger.warning("Failure in HALF_OPEN state, transitioning back to OPEN")
+                self._state = CircuitState.OPEN
+                self._success_count = 0
+                self._last_failure_time = time.time()

code_puppy/mcp/config_wizard.py

@@ -0,0 +1,437 @@
+"""
+MCP Configuration Wizard - Interactive setup for MCP servers.
+"""
+
+import re
+from typing import Dict, Optional
+from urllib.parse import urlparse
+
+from code_puppy.mcp import ServerConfig, get_mcp_manager
+from code_puppy.messaging import emit_error, emit_info, emit_success, emit_warning
+from rich.prompt import Prompt, Confirm
+from rich.console import Console
+
+console = Console()
+
+
+class MCPConfigWizard:
+    """Interactive wizard for configuring MCP servers."""
+    
+    def __init__(self):
+        self.manager = get_mcp_manager()
+    
+    def run_wizard(self) -> Optional[ServerConfig]:
+        """
+        Run the interactive configuration wizard.
+        
+        Returns:
+            ServerConfig if successful, None if cancelled
+        """
+        console.print("\n[bold cyan]🧙 MCP Server Configuration Wizard[/bold cyan]\n")
+        
+        # Step 1: Server name
+        name = self.prompt_server_name()
+        if not name:
+            return None
+        
+        # Step 2: Server type
+        server_type = self.prompt_server_type()
+        if not server_type:
+            return None
+        
+        # Step 3: Type-specific configuration
+        config = {}
+        if server_type == "sse":
+            config = self.prompt_sse_config()
+        elif server_type == "http":
+            config = self.prompt_http_config()
+        elif server_type == "stdio":
+            config = self.prompt_stdio_config()
+        
+        if not config:
+            return None
+        
+        # Step 4: Create ServerConfig
+        server_config = ServerConfig(
+            id=f"{name}_{hash(name)}",
+            name=name,
+            type=server_type,
+            enabled=True,
+            config=config
+        )
+        
+        # Step 5: Show summary and confirm
+        if self.prompt_confirmation(server_config):
+            return server_config
+        
+        return None
+    
+    def prompt_server_name(self) -> Optional[str]:
+        """Prompt for server name with validation."""
+        while True:
+            name = Prompt.ask(
+                "[yellow]Enter server name[/yellow]",
+                default=None
+            )
+            
+            if not name:
+                if not Confirm.ask("Cancel configuration?", default=False):
+                    continue
+                return None
+            
+            # Validate name
+            if not self.validate_name(name):
+                emit_error("Name must be alphanumeric with hyphens/underscores only")
+                continue
+            
+            # Check uniqueness
+            existing = self.manager.registry.get_by_name(name)
+            if existing:
+                emit_error(f"Server '{name}' already exists")
+                continue
+            
+            return name
+    
+    def prompt_server_type(self) -> Optional[str]:
+        """Prompt for server type."""
+        console.print("\n[cyan]Server types:[/cyan]")
+        console.print("  [bold]sse[/bold]   - Server-Sent Events (HTTP streaming)")
+        console.print("  [bold]http[/bold]  - HTTP/REST API")
+        console.print("  [bold]stdio[/bold] - Local command (subprocess)")
+        
+        while True:
+            server_type = Prompt.ask(
+                "\n[yellow]Select server type[/yellow]",
+                choices=["sse", "http", "stdio"],
+                default="stdio"
+            )
+            
+            if server_type in ["sse", "http", "stdio"]:
+                return server_type
+            
+            emit_error("Invalid type. Choose: sse, http, or stdio")
+    
+    def prompt_sse_config(self) -> Optional[Dict]:
+        """Prompt for SSE server configuration."""
+        console.print("\n[cyan]Configuring SSE server[/cyan]")
+        
+        # URL
+        url = self.prompt_url("SSE")
+        if not url:
+            return None
+        
+        config = {
+            "type": "sse",
+            "url": url,
+            "timeout": 30
+        }
+        
+        # Headers (optional)
+        if Confirm.ask("Add custom headers?", default=False):
+            headers = self.prompt_headers()
+            if headers:
+                config["headers"] = headers
+        
+        # Timeout
+        timeout_str = Prompt.ask(
+            "Connection timeout (seconds)",
+            default="30"
+        )
+        try:
+            config["timeout"] = int(timeout_str)
+        except ValueError:
+            config["timeout"] = 30
+        
+        return config
+    
+    def prompt_http_config(self) -> Optional[Dict]:
+        """Prompt for HTTP server configuration."""
+        console.print("\n[cyan]Configuring HTTP server[/cyan]")
+        
+        # URL
+        url = self.prompt_url("HTTP")
+        if not url:
+            return None
+        
+        config = {
+            "type": "http",
+            "url": url,
+            "timeout": 30
+        }
+        
+        # Headers (optional)
+        if Confirm.ask("Add custom headers?", default=False):
+            headers = self.prompt_headers()
+            if headers:
+                config["headers"] = headers
+        
+        # Timeout
+        timeout_str = Prompt.ask(
+            "Request timeout (seconds)",
+            default="30"
+        )
+        try:
+            config["timeout"] = int(timeout_str)
+        except ValueError:
+            config["timeout"] = 30
+        
+        return config
+    
+    def prompt_stdio_config(self) -> Optional[Dict]:
+        """Prompt for Stdio server configuration."""
+        console.print("\n[cyan]Configuring Stdio server[/cyan]")
+        console.print("[dim]Examples:[/dim]")
+        console.print("[dim]  • npx -y @modelcontextprotocol/server-filesystem /path[/dim]")
+        console.print("[dim]  • python mcp_server.py[/dim]")
+        console.print("[dim]  • node server.js[/dim]")
+        
+        # Command
+        command = Prompt.ask(
+            "\n[yellow]Enter command[/yellow]",
+            default=None
+        )
+        
+        if not command:
+            return None
+        
+        config = {
+            "type": "stdio",
+            "command": command,
+            "args": [],
+            "timeout": 30
+        }
+        
+        # Arguments
+        args_str = Prompt.ask(
+            "Enter arguments (space-separated)",
+            default=""
+        )
+        if args_str:
+            # Simple argument parsing (handles quoted strings)
+            import shlex
+            try:
+                config["args"] = shlex.split(args_str)
+            except ValueError:
+                config["args"] = args_str.split()
+        
+        # Working directory (optional)
+        cwd = Prompt.ask(
+            "Working directory (optional)",
+            default=""
+        )
+        if cwd:
+            import os
+            if os.path.isdir(os.path.expanduser(cwd)):
+                config["cwd"] = os.path.expanduser(cwd)
+            else:
+                emit_warning(f"Directory '{cwd}' not found, ignoring")
+        
+        # Environment variables (optional)
+        if Confirm.ask("Add environment variables?", default=False):
+            env = self.prompt_env_vars()
+            if env:
+                config["env"] = env
+        
+        # Timeout
+        timeout_str = Prompt.ask(
+            "Startup timeout (seconds)",
+            default="30"
+        )
+        try:
+            config["timeout"] = int(timeout_str)
+        except ValueError:
+            config["timeout"] = 30
+        
+        return config
+    
+    def prompt_url(self, server_type: str) -> Optional[str]:
+        """Prompt for and validate URL."""
+        while True:
+            url = Prompt.ask(
+                f"[yellow]Enter {server_type} server URL[/yellow]",
+                default=None
+            )
+            
+            if not url:
+                if Confirm.ask("Cancel configuration?", default=False):
+                    return None
+                continue
+            
+            if self.validate_url(url):
+                return url
+            
+            emit_error("Invalid URL. Must be http:// or https://")
+    
+    def prompt_headers(self) -> Dict[str, str]:
+        """Prompt for HTTP headers."""
+        headers = {}
+        console.print("[dim]Enter headers (format: Name: Value)[/dim]")
+        console.print("[dim]Press Enter with empty name to finish[/dim]")
+        
+        while True:
+            name = Prompt.ask("Header name", default="")
+            if not name:
+                break
+            
+            value = Prompt.ask(f"Value for '{name}'", default="")
+            headers[name] = value
+            
+            if not Confirm.ask("Add another header?", default=True):
+                break
+        
+        return headers
+    
+    def prompt_env_vars(self) -> Dict[str, str]:
+        """Prompt for environment variables."""
+        env = {}
+        console.print("[dim]Enter environment variables[/dim]")
+        console.print("[dim]Press Enter with empty name to finish[/dim]")
+        
+        while True:
+            name = Prompt.ask("Variable name", default="")
+            if not name:
+                break
+            
+            value = Prompt.ask(f"Value for '{name}'", default="")
+            env[name] = value
+            
+            if not Confirm.ask("Add another variable?", default=True):
+                break
+        
+        return env
+    
+    def validate_name(self, name: str) -> bool:
+        """Validate server name."""
+        # Allow alphanumeric, hyphens, and underscores
+        return bool(re.match(r'^[a-zA-Z0-9_-]+$', name))
+    
+    def validate_url(self, url: str) -> bool:
+        """Validate URL format."""
+        try:
+            result = urlparse(url)
+            return result.scheme in ('http', 'https') and bool(result.netloc)
+        except Exception:
+            return False
+    
+    def validate_command(self, command: str) -> bool:
+        """Check if command exists (basic check)."""
+        import shutil
+        import os
+        
+        # If it's a path, check if file exists
+        if '/' in command or '\\' in command:
+            return os.path.isfile(command)
+        
+        # Otherwise check if it's in PATH
+        return shutil.which(command) is not None
+    
+    def test_connection(self, config: ServerConfig) -> bool:
+        """
+        Test connection to the configured server.
+        
+        Args:
+            config: Server configuration to test
+            
+        Returns:
+            True if connection successful, False otherwise
+        """
+        emit_info("Testing connection...")
+        
+        try:
+            # Try to create the server instance
+            managed = self.manager.get_server(config.id)
+            if not managed:
+                # Temporarily register to test
+                self.manager.register_server(config)
+                managed = self.manager.get_server(config.id)
+            
+            if managed:
+                # Try to get the pydantic server (this validates config)
+                server = managed.get_pydantic_server()
+                if server:
+                    emit_success("✓ Configuration valid")
+                    return True
+            
+            emit_error("✗ Failed to create server instance")
+            return False
+            
+        except Exception as e:
+            emit_error(f"✗ Configuration error: {e}")
+            return False
+    
+    def prompt_confirmation(self, config: ServerConfig) -> bool:
+        """Show summary and ask for confirmation."""
+        console.print("\n[bold cyan]Configuration Summary:[/bold cyan]")
+        console.print(f"  [bold]Name:[/bold] {config.name}")
+        console.print(f"  [bold]Type:[/bold] {config.type}")
+        
+        if config.type in ["sse", "http"]:
+            console.print(f"  [bold]URL:[/bold] {config.config.get('url')}")
+        elif config.type == "stdio":
+            console.print(f"  [bold]Command:[/bold] {config.config.get('command')}")
+            args = config.config.get('args', [])
+            if args:
+                console.print(f"  [bold]Arguments:[/bold] {' '.join(args)}")
+        
+        console.print(f"  [bold]Timeout:[/bold] {config.config.get('timeout', 30)}s")
+        
+        # Test connection if requested
+        if Confirm.ask("\n[yellow]Test connection?[/yellow]", default=True):
+            if not self.test_connection(config):
+                if not Confirm.ask("Continue anyway?", default=False):
+                    return False
+        
+        return Confirm.ask("\n[bold green]Save this configuration?[/bold green]", default=True)
+
+
+def run_add_wizard() -> bool:
+    """
+    Run the MCP add wizard and register the server.
+    
+    Returns:
+        True if server was added, False otherwise
+    """
+    wizard = MCPConfigWizard()
+    config = wizard.run_wizard()
+    
+    if config:
+        try:
+            manager = get_mcp_manager()
+            server_id = manager.register_server(config)
+            
+            emit_success(f"\n✅ Server '{config.name}' added successfully!")
+            emit_info(f"Server ID: {server_id}")
+            emit_info("Use '/mcp list' to see all servers")
+            emit_info(f"Use '/mcp start {config.name}' to start the server")
+            
+            # Also save to mcp_servers.json for persistence
+            from code_puppy.config import MCP_SERVERS_FILE, load_mcp_server_configs
+            import json
+            import os
+            
+            # Load existing configs
+            if os.path.exists(MCP_SERVERS_FILE):
+                with open(MCP_SERVERS_FILE, 'r') as f:
+                    data = json.load(f)
+                    servers = data.get("mcp_servers", {})
+            else:
+                servers = {}
+                data = {"mcp_servers": servers}
+            
+            # Add new server
+            servers[config.name] = config.config
+            
+            # Save back
+            os.makedirs(os.path.dirname(MCP_SERVERS_FILE), exist_ok=True)
+            with open(MCP_SERVERS_FILE, 'w') as f:
+                json.dump(data, f, indent=2)
+            
+            emit_info(f"[dim]Configuration saved to {MCP_SERVERS_FILE}[/dim]")
+            return True
+            
+        except Exception as e:
+            emit_error(f"Failed to add server: {e}")
+            return False
+    else:
+        emit_warning("Configuration cancelled")
+        return False