mseep-lightfast-mcp 0.0.1__py3-none-any.whl

tools/orchestration/server_orchestrator.py

@@ -0,0 +1,466 @@
+"""Server orchestrator for managing multiple MCP servers simultaneously."""
+
+import signal
+import subprocess
+import threading
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Dict, List, Optional
+
+from lightfast_mcp.core.base_server import BaseServer, ServerConfig
+from tools.common import (
+    OperationStatus,
+    Result,
+    RetryManager,
+    ServerInfo,
+    ServerStartupError,
+    ServerState,
+    get_logger,
+    run_concurrent_operations,
+    with_correlation_id,
+    with_operation_context,
+)
+
+from .server_registry import get_registry
+
+logger = get_logger("ServerOrchestrator")
+
+
+@dataclass
+class ServerProcess:
+    """Information about a running server process."""
+
+    server: Optional[BaseServer] = None
+    thread: Optional[threading.Thread] = None
+    process: Optional[subprocess.Popen] = None
+    process_id: Optional[int] = None
+    start_time: datetime = field(default_factory=datetime.utcnow)
+    is_background: bool = False
+    config: Optional[ServerConfig] = None
+
+    @property
+    def uptime_seconds(self) -> float:
+        return (datetime.utcnow() - self.start_time).total_seconds()
+
+
+class ServerOrchestrator:
+    """Orchestrates lifecycle and coordination of multiple MCP servers."""
+
+    def __init__(self, max_concurrent_startups: int = 3):
+        """Initialize the server orchestrator."""
+        self.registry = get_registry()
+        self._running_servers: Dict[str, ServerProcess] = {}
+        self._shutdown_event = threading.Event()
+        self.max_concurrent_startups = max_concurrent_startups
+        self.retry_manager = RetryManager(max_attempts=3, base_delay=2.0)
+
+        # Setup signal handlers for graceful shutdown
+        self._setup_signal_handlers()
+
+    def _setup_signal_handlers(self):
+        """Setup signal handlers for graceful shutdown."""
+
+        def signal_handler(signum, frame):
+            logger.info(f"Received signal {signum}, initiating graceful shutdown")
+            self.shutdown_all()
+
+        try:
+            signal.signal(signal.SIGINT, signal_handler)
+            signal.signal(signal.SIGTERM, signal_handler)
+        except ValueError:
+            # Signal handling might not work in some environments (like Jupyter)
+            logger.debug("Could not setup signal handlers")
+
+    @with_correlation_id
+    @with_operation_context(operation="start_server")
+    async def start_server(
+        self,
+        server_config: ServerConfig,
+        background: bool = False,
+        show_logs: bool = True,
+    ) -> Result[ServerInfo]:
+        """Start a single server with proper error handling."""
+        if server_config.name in self._running_servers:
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"Server {server_config.name} is already running",
+                error_code="SERVER_ALREADY_RUNNING",
+            )
+
+        # Validate configuration
+        server_type = server_config.config.get("type", "unknown")
+        is_valid, error_msg = self.registry.validate_server_config(
+            server_type, server_config
+        )
+        if not is_valid:
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"Invalid configuration: {error_msg}",
+                error_code="INVALID_CONFIG",
+            )
+
+        try:
+            # Choose startup method based on transport
+            if server_config.transport in ["http", "streamable-http"]:
+                result = await self._start_server_subprocess(
+                    server_config, background, show_logs
+                )
+            else:
+                result = await self._start_server_traditional(server_config, background)
+
+            if result.is_success:
+                logger.info(
+                    f"Successfully started server: {server_config.name}",
+                    server_name=server_config.name,
+                    transport=server_config.transport,
+                    background=background,
+                )
+
+            return result
+
+        except Exception as e:
+            error = ServerStartupError(
+                f"Failed to start server {server_config.name}",
+                server_name=server_config.name,
+                cause=e,
+            )
+            logger.error("Server startup failed", error=error)
+            return Result(
+                status=OperationStatus.FAILED,
+                error=str(error),
+                error_code=error.error_code,
+            )
+
+    async def _start_server_subprocess(
+        self, server_config: ServerConfig, background: bool, show_logs: bool = True
+    ) -> Result[ServerInfo]:
+        """Start a server using subprocess (for HTTP transports)."""
+        try:
+            # Determine the correct script based on server type
+            server_type = server_config.config.get("type", "unknown")
+
+            # Validated list of trusted server modules (security: only allow known modules)
+            trusted_modules = {
+                "mock": "lightfast_mcp.servers.mock_server",
+                "blender": "lightfast_mcp.servers.blender_mcp_server",
+            }
+
+            if server_type not in trusted_modules:
+                return Result(
+                    status=OperationStatus.FAILED,
+                    error=f"Unknown or untrusted server type for subprocess: {server_type}",
+                    error_code="UNTRUSTED_SERVER_TYPE",
+                )
+
+            script_module = trusted_modules[server_type]
+
+            # Create environment with server config
+            import json
+            import os
+            import shutil
+            import subprocess
+            import sys
+
+            env = os.environ.copy()
+            env["LIGHTFAST_MCP_SERVER_CONFIG"] = json.dumps(
+                {
+                    "name": server_config.name,
+                    "description": server_config.description,
+                    "host": server_config.host,
+                    "port": server_config.port,
+                    "transport": server_config.transport,
+                    "path": server_config.path,
+                    "config": server_config.config,
+                }
+            )
+
+            # Start the subprocess
+            # Security: Use full path to python executable and validated module name
+            python_executable = (
+                sys.executable or shutil.which("python3") or shutil.which("python")
+            )
+            if not python_executable:
+                return Result(
+                    status=OperationStatus.FAILED,
+                    error="Could not find Python executable",
+                    error_code="PYTHON_NOT_FOUND",
+                )
+
+            # Only capture logs if background=True AND show_logs=False
+            capture_logs = background and not show_logs
+            process = subprocess.Popen(
+                [python_executable, "-m", script_module],  # nosec B603 - using validated module and full path
+                env=env,
+                stdout=subprocess.PIPE if capture_logs else None,
+                stderr=subprocess.PIPE if capture_logs else None,
+                text=True,
+            )
+
+            server_process = ServerProcess(
+                process=process,
+                process_id=process.pid,
+                is_background=background,
+                config=server_config,
+            )
+
+            self._running_servers[server_config.name] = server_process
+
+            # Create ServerInfo
+            server_info = ServerInfo(
+                name=server_config.name,
+                server_type=server_type,
+                state=ServerState.RUNNING,
+                host=server_config.host,
+                port=server_config.port,
+                transport=server_config.transport,
+                url=f"http://{server_config.host}:{server_config.port}{server_config.path}",
+                pid=process.pid,
+                start_time=server_process.start_time,
+            )
+
+            return Result(status=OperationStatus.SUCCESS, data=server_info)
+
+        except Exception as e:
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"Failed to start server subprocess: {e}",
+                error_code="SUBPROCESS_START_FAILED",
+            )
+
+    async def _start_server_traditional(
+        self, server_config: ServerConfig, background: bool
+    ) -> Result[ServerInfo]:
+        """Start a server using the traditional approach (for stdio transport)."""
+        try:
+            # Create server instance
+            server_type = server_config.config.get("type", "unknown")
+            if not server_type or server_type == "unknown":
+                return Result(
+                    status=OperationStatus.FAILED,
+                    error="Server type not specified in configuration",
+                    error_code="MISSING_SERVER_TYPE",
+                )
+            server = self.registry.create_server(server_type, server_config)
+
+            if background:
+                # Run in background thread with proper asyncio handling
+                import threading
+
+                thread = threading.Thread(
+                    target=self._run_server_in_thread,
+                    args=(server, server_config.name),
+                    daemon=True,
+                )
+                thread.start()
+
+                server_process = ServerProcess(
+                    server=server,
+                    thread=thread,
+                    is_background=True,
+                    config=server_config,
+                )
+            else:
+                # Run in foreground (blocking)
+                server_process = ServerProcess(
+                    server=server, is_background=False, config=server_config
+                )
+                # This will block
+                server.run()
+
+            self._running_servers[server_config.name] = server_process
+
+            # Create ServerInfo
+            server_info = ServerInfo(
+                name=server_config.name,
+                server_type=server_config.config.get("type", "unknown"),
+                state=ServerState.RUNNING,
+                host=server_config.host,
+                port=server_config.port,
+                transport=server_config.transport,
+                start_time=server_process.start_time,
+            )
+
+            return Result(status=OperationStatus.SUCCESS, data=server_info)
+
+        except Exception as e:
+            return Result(
+                status=OperationStatus.FAILED,
+                error=f"Failed to start server traditionally: {e}",
+                error_code="TRADITIONAL_START_FAILED",
+            )
+
+    def _run_server_in_thread(self, server: BaseServer, server_name: str):
+        """Run a server in a background thread with proper asyncio handling."""
+        try:
+            logger.info(f"Running server {server_name} in background thread")
+
+            # Create a new event loop for this thread
+            # This is crucial for FastMCP servers which are asyncio-based
+            import asyncio
+
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+            try:
+                # Run the server in the new event loop
+                server.run()
+            finally:
+                # Clean up the event loop
+                loop.close()
+
+        except Exception as e:
+            logger.error(f"Error running server {server_name}", error=e)
+        finally:
+            # Clean up the server from running list when it stops
+            if server_name in self._running_servers:
+                logger.info(f"Server {server_name} stopped")
+                del self._running_servers[server_name]
+
+    @with_correlation_id
+    async def start_multiple_servers(
+        self,
+        server_configs: List[ServerConfig],
+        background: bool = True,
+        show_logs: bool = True,
+    ) -> Result[Dict[str, bool]]:
+        """Start multiple servers concurrently."""
+        if not server_configs:
+            return Result(status=OperationStatus.SUCCESS, data={})
+
+        logger.info(f"Starting {len(server_configs)} servers concurrently")
+
+        # Create startup operations
+        async def start_single_server(config: ServerConfig) -> bool:
+            result = await self.start_server(config, background, show_logs)
+            return result.is_success
+
+        operations = [
+            lambda cfg=config: start_single_server(cfg) for config in server_configs
+        ]
+        operation_names = [f"start_{config.name}" for config in server_configs]
+
+        # Execute with controlled concurrency
+        results = await run_concurrent_operations(
+            operations,
+            max_concurrent=self.max_concurrent_startups,
+            operation_names=operation_names,
+        )
+
+        # Build result dictionary
+        startup_results = {}
+        for config, result in zip(server_configs, results):
+            # If the operation succeeded, use the data (boolean result)
+            # If the operation failed (exception), treat as False
+            startup_results[config.name] = result.data if result.is_success else False
+
+        successful = sum(1 for success in startup_results.values() if success)
+        logger.info(
+            f"Server startup completed: {successful}/{len(server_configs)} successful"
+        )
+
+        return Result(status=OperationStatus.SUCCESS, data=startup_results)
+
+    def get_running_servers(self) -> Dict[str, ServerInfo]:
+        """Get information about all running servers."""
+        info = {}
+        for name, server_process in self._running_servers.items():
+            if server_process.server:
+                info[name] = server_process.server.info
+            elif server_process.config:
+                # Create ServerInfo for subprocess-based servers
+                server_info = ServerInfo(
+                    name=server_process.config.name,
+                    server_type=server_process.config.config.get("type", "unknown"),
+                    state=ServerState.RUNNING
+                    if self._is_process_running(server_process)
+                    else ServerState.ERROR,
+                    host=server_process.config.host,
+                    port=server_process.config.port,
+                    transport=server_process.config.transport,
+                    start_time=server_process.start_time,
+                    pid=server_process.process_id,
+                )
+                if server_process.config.transport in ["http", "streamable-http"]:
+                    server_info.url = (
+                        f"http://{server_process.config.host}:{server_process.config.port}"
+                        f"{server_process.config.path}"
+                    )
+                info[name] = server_info
+        return info
+
+    def _is_process_running(self, server_process: ServerProcess) -> bool:
+        """Check if a subprocess is still running."""
+        if server_process.process:
+            return server_process.process.poll() is None
+        return False
+
+    def shutdown_all(self):
+        """Shutdown all running servers."""
+        logger.info("Shutting down all servers...")
+
+        server_names = list(self._running_servers.keys())
+        for name in server_names:
+            self.stop_server(name)
+
+        self._shutdown_event.set()
+        logger.info("All servers shut down")
+
+    def stop_server(self, server_name: str) -> bool:
+        """Stop a specific server."""
+        if server_name not in self._running_servers:
+            logger.warning(f"Server {server_name} is not running")
+            return False
+
+        server_process = self._running_servers[server_name]
+
+        try:
+            # Handle subprocess
+            if server_process.process:
+                try:
+                    server_process.process.terminate()
+                    # Wait up to 5 seconds for graceful shutdown
+                    server_process.process.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    logger.warning(
+                        f"Server {server_name} didn't stop gracefully, forcing..."
+                    )
+                    server_process.process.kill()
+                except Exception as e:
+                    logger.error(
+                        f"Error stopping server process {server_name}", error=e
+                    )
+
+            # Handle threaded server
+            elif server_process.thread and server_process.thread.is_alive():
+                logger.info(
+                    f"Server {server_name} thread is still running, waiting for natural shutdown..."
+                )
+
+            # Remove from running servers
+            del self._running_servers[server_name]
+            logger.info(f"Stopped server: {server_name}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Error stopping server {server_name}", error=e)
+            return False
+
+
+# Global orchestrator instance
+_orchestrator: Optional[ServerOrchestrator] = None
+
+
+def get_orchestrator() -> ServerOrchestrator:
+    """Get the global server orchestrator instance."""
+    global _orchestrator
+    if _orchestrator is None:
+        _orchestrator = ServerOrchestrator()
+    return _orchestrator
+
+
+def reset_orchestrator():
+    """Reset the global orchestrator (mainly for testing)."""
+    global _orchestrator
+    if _orchestrator is not None:
+        _orchestrator.shutdown_all()
+    _orchestrator = None

tools/orchestration/server_registry.py

@@ -0,0 +1,187 @@
+"""Server registry for discovering and managing MCP servers."""
+
+import importlib
+import inspect
+import pkgutil
+from pathlib import Path
+from typing import Any
+
+from lightfast_mcp.core.base_server import BaseServer, ServerConfig
+from lightfast_mcp.utils.logging_utils import get_logger
+
+logger = get_logger("ServerRegistry")
+
+
+class ServerRegistry:
+    """Registry for discovering and managing MCP servers."""
+
+    def __init__(self):
+        """Initialize the server registry."""
+        self._server_classes: dict[str, type[BaseServer]] = {}
+        self._server_instances: dict[str, BaseServer] = {}
+
+        # Auto-discover servers on initialization
+        self.discover_servers()
+
+    def discover_servers(self):
+        """Automatically discover all available server classes."""
+        logger.info("Discovering MCP servers...")
+
+        # Discover servers in the servers package
+        self._discover_servers_in_package("lightfast_mcp.servers")
+
+        logger.info(
+            f"Found {len(self._server_classes)} server types: {list(self._server_classes.keys())}"
+        )
+
+    def _discover_servers_in_package(self, package_name: str):
+        """Discover servers in a specific package."""
+        try:
+            # Import the package
+            package = importlib.import_module(package_name)
+
+            # Check if package has a file path
+            if package.__file__ is None:
+                logger.debug(f"Package {package_name} has no __file__ attribute")
+                return
+
+            package_path = Path(package.__file__).parent
+
+            # Iterate through all modules in the package
+            for _finder, name, ispkg in pkgutil.iter_modules([str(package_path)]):
+                if ispkg:
+                    # Recursively search subpackages
+                    subpackage_name = f"{package_name}.{name}"
+                    self._discover_servers_in_package(subpackage_name)
+                else:
+                    # Import the module and look for server classes
+                    try:
+                        module_name = f"{package_name}.{name}"
+                        module = importlib.import_module(module_name)
+                        self._discover_servers_in_module(module)
+                    except Exception as e:
+                        logger.debug(f"Could not import module {module_name}: {e}")
+
+        except Exception as e:
+            logger.debug(f"Could not discover servers in package {package_name}: {e}")
+
+    def _discover_servers_in_module(self, module):
+        """Discover server classes in a module."""
+        for name, obj in inspect.getmembers(module):
+            if (
+                inspect.isclass(obj)
+                and issubclass(obj, BaseServer)
+                and obj is not BaseServer
+                and hasattr(obj, "SERVER_TYPE")
+            ):
+                server_type = obj.SERVER_TYPE
+                if server_type != "base":  # Skip the base class
+                    logger.debug(f"Found server class: {name} (type: {server_type})")
+                    self._server_classes[server_type] = obj
+
+    def register_server_class(self, server_type: str, server_class: type[BaseServer]):
+        """Manually register a server class."""
+        logger.info(f"Registering server class: {server_type}")
+        self._server_classes[server_type] = server_class
+
+    def get_server_class(self, server_type: str) -> type[BaseServer] | None:
+        """Get a server class by type."""
+        return self._server_classes.get(server_type)
+
+    def get_available_server_types(self) -> list[str]:
+        """Get list of all available server types."""
+        return list(self._server_classes.keys())
+
+    def create_server(self, server_type: str, config: ServerConfig) -> BaseServer:
+        """Create a server instance from configuration."""
+        server_class = self.get_server_class(server_type)
+        if not server_class:
+            raise ValueError(f"Unknown server type: {server_type}")
+
+        logger.info(f"Creating server instance: {config.name} (type: {server_type})")
+        server = server_class.create_from_config(config)
+
+        # Store the instance
+        self._server_instances[config.name] = server
+        return server
+
+    def get_server_instance(self, name: str) -> BaseServer | None:
+        """Get a running server instance by name."""
+        return self._server_instances.get(name)
+
+    def get_all_server_instances(self) -> dict[str, BaseServer]:
+        """Get all server instances."""
+        return self._server_instances.copy()
+
+    def remove_server_instance(self, name: str) -> bool:
+        """Remove a server instance."""
+        if name in self._server_instances:
+            logger.info(f"Removing server instance: {name}")
+            del self._server_instances[name]
+            return True
+        return False
+
+    def get_server_info(self) -> dict[str, dict[str, Any]]:
+        """Get information about all available server types."""
+        info = {}
+        for server_type, server_class in self._server_classes.items():
+            info[server_type] = {
+                "class_name": server_class.__name__,
+                "version": getattr(server_class, "SERVER_VERSION", "1.0.0"),
+                "required_dependencies": getattr(
+                    server_class, "REQUIRED_DEPENDENCIES", []
+                ),
+                "required_apps": getattr(server_class, "REQUIRED_APPS", []),
+                "description": getattr(
+                    server_class, "__doc__", "No description available"
+                ),
+            }
+        return info
+
+    def validate_server_config(
+        self, server_type: str, config: ServerConfig
+    ) -> tuple[bool, str]:
+        """Validate a server configuration."""
+        server_class = self.get_server_class(server_type)
+        if not server_class:
+            return False, f"Unknown server type: {server_type}"
+
+        # Basic validation
+        if not config.name:
+            return False, "Server name is required"
+
+        if not config.description:
+            return False, "Server description is required"
+
+        # Check for port conflicts with existing instances
+        for instance in self._server_instances.values():
+            if (
+                instance.config.port == config.port
+                and instance.config.host == config.host
+                and instance.config.transport in ["http", "streamable-http"]
+                and config.transport in ["http", "streamable-http"]
+            ):
+                return (
+                    False,
+                    f"Port {config.port} on {config.host} is already in use by {instance.config.name}",
+                )
+
+        return True, "Configuration is valid"
+
+
+# Global registry instance
+_registry: ServerRegistry | None = None
+
+
+def get_registry() -> ServerRegistry:
+    """Get the global server registry instance."""
+    global _registry
+    if _registry is None:
+        _registry = ServerRegistry()
+    return _registry
+
+
+def reset_registry():
+    """Reset the global registry (mainly for testing)."""
+    global _registry
+    _registry = None