stravinsky 0.2.67__py3-none-any.whl → 0.4.18__py3-none-any.whl

mcp_bridge/tools/lsp/manager.py

@@ -0,0 +1,448 @@
+"""
+Persistent LSP Server Manager
+
+Manages persistent Language Server Protocol (LSP) servers for improved performance.
+Implements lazy initialization, JSON-RPC communication, and graceful shutdown.
+
+Architecture:
+- Servers start on first use (lazy initialization)
+- JSON-RPC over stdio using pygls BaseLanguageClient
+- Supports Python (jedi-language-server) and TypeScript (typescript-language-server)
+- Graceful shutdown on MCP server exit
+- Health checks and idle timeout management
+"""
+
+import asyncio
+import json
+import logging
+import os
+import shlex
+import signal
+import threading
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Optional
+
+from pygls.client import JsonRPCClient
+from lsprotocol.types import (
+    InitializeParams,
+    InitializedParams,
+    ClientCapabilities,
+    WorkspaceFolder,
+)
+
+logger = logging.getLogger(__name__)
+
+# Configuration for LSP server lifecycle management
+LSP_CONFIG = {
+    "idle_timeout": 1800,  # 30 minutes
+    "health_check_interval": 300,  # 5 minutes
+    "health_check_timeout": 5.0,
+}
+
+
+@dataclass
+class LSPServer:
+    """Metadata for a persistent LSP server."""
+
+    name: str
+    command: list[str]
+    client: Optional[JsonRPCClient] = None
+    initialized: bool = False
+    process: Optional[asyncio.subprocess.Process] = None
+    pid: Optional[int] = None  # Track subprocess PID for explicit cleanup
+    last_used: float = field(default_factory=time.time)  # Timestamp of last usage
+    created_at: float = field(default_factory=time.time)  # Timestamp of server creation
+
+
+class LSPManager:
+    """
+    Singleton manager for persistent LSP servers.
+
+    Implements:
+    - Lazy server initialization (start on first use)
+    - Process lifecycle management with GC protection
+    - Exponential backoff for crash recovery
+    - Graceful shutdown with signal handling
+    - Health checks and idle server shutdown
+    """
+
+    _instance: Optional["LSPManager"] = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        if hasattr(self, "_initialized"):
+            return
+        self._initialized = True
+        self._servers: dict[str, LSPServer] = {}
+        self._lock = asyncio.Lock()
+        self._restart_attempts: dict[str, int] = {}
+        self._health_monitor_task: Optional[asyncio.Task] = None
+
+        # Register available LSP servers
+        self._register_servers()
+
+    def _register_servers(self):
+        """Register available LSP server configurations."""
+        self._servers["python"] = LSPServer(name="python", command=["jedi-language-server"])
+        self._servers["typescript"] = LSPServer(
+            name="typescript", command=["typescript-language-server", "--stdio"]
+        )
+
+    async def get_server(self, language: str) -> Optional[JsonRPCClient]:
+        """
+        Get or start a persistent LSP server for the given language.
+
+        Args:
+            language: Language identifier (e.g., "python", "typescript")
+
+        Returns:
+            JsonRPCClient instance or None if server unavailable
+        """
+        if language not in self._servers:
+            logger.warning(f"No LSP server configured for language: {language}")
+            return None
+
+        server = self._servers[language]
+
+        # Return existing initialized server
+        if server.initialized and server.client:
+            # Update last_used timestamp
+            server.last_used = time.time()
+            # Start health monitor on first use
+            if self._health_monitor_task is None or self._health_monitor_task.done():
+                self._health_monitor_task = asyncio.create_task(self._background_health_monitor())
+            return server.client
+
+        # Start server with lock to prevent race conditions
+        async with self._lock:
+            # Double-check after acquiring lock
+            if server.initialized and server.client:
+                server.last_used = time.time()
+                if self._health_monitor_task is None or self._health_monitor_task.done():
+                    self._health_monitor_task = asyncio.create_task(self._background_health_monitor())
+                return server.client
+
+            try:
+                await self._start_server(server)
+                # Start health monitor on first server creation
+                if self._health_monitor_task is None or self._health_monitor_task.done():
+                    self._health_monitor_task = asyncio.create_task(self._background_health_monitor())
+                return server.client
+            except Exception as e:
+                logger.error(f"Failed to start {language} LSP server: {e}")
+                return None
+
+    async def _start_server(self, server: LSPServer):
+        """
+        Start a persistent LSP server process.
+
+        Implements:
+        - Process health validation after start
+        - LSP initialization handshake
+        - GC protection via persistent reference
+        - Timestamp tracking for idle detection
+
+        Args:
+            server: LSPServer metadata object
+        """
+        try:
+            # Create pygls client
+            client = JsonRPCClient()
+
+            logger.info(f"Starting {server.name} LSP server: {' '.join(server.command)}")
+
+            # Start server process (start_io expects cmd as first arg, then *args)
+            await client.start_io(server.command[0], *server.command[1:])
+
+            # Brief delay for process startup
+            await asyncio.sleep(0.2)
+
+            # Capture subprocess from client
+            if not hasattr(client, "_server") or client._server is None:
+                raise ConnectionError(
+                    f"{server.name} LSP server process not accessible after start_io()"
+                )
+
+            server.process = client._server
+            server.pid = server.process.pid
+            logger.debug(f"{server.name} LSP server started with PID: {server.pid}")
+
+            # Validate process is still running
+            if server.process.returncode is not None:
+                raise ConnectionError(
+                    f"{server.name} LSP server exited immediately (code {server.process.returncode})"
+                )
+
+            # Perform LSP initialization handshake
+            init_params = InitializeParams(
+                process_id=None, root_uri=None, capabilities=ClientCapabilities()
+            )
+
+            try:
+                # Send initialize request via protocol
+                response = await asyncio.wait_for(
+                    client.protocol.send_request_async("initialize", init_params), timeout=10.0
+                )
+
+                # Send initialized notification
+                client.protocol.notify("initialized", InitializedParams())
+
+                logger.info(f"{server.name} LSP server initialized: {response}")
+
+            except asyncio.TimeoutError:
+                raise ConnectionError(f"{server.name} LSP server initialization timed out")
+
+            # Store client reference (GC protection)
+            server.client = client
+            server.initialized = True
+            server.created_at = time.time()
+            server.last_used = time.time()
+
+            # Reset restart attempts on successful start
+            self._restart_attempts[server.name] = 0
+
+            logger.info(f"{server.name} LSP server started successfully")
+
+        except Exception as e:
+            logger.error(f"Failed to start {server.name} LSP server: {e}", exc_info=True)
+            # Cleanup on failure
+            if server.client:
+                try:
+                    await server.client.stop()
+                except:
+                    pass
+            server.client = None
+            server.initialized = False
+            server.process = None
+            server.pid = None
+            raise
+
+    async def _restart_with_backoff(self, server: LSPServer):
+        """
+        Restart a crashed LSP server with exponential backoff.
+
+        Strategy: delay = 2^attempt + jitter (max 60s)
+
+        Args:
+            server: LSPServer to restart
+        """
+        import random
+
+        attempt = self._restart_attempts.get(server.name, 0)
+        self._restart_attempts[server.name] = attempt + 1
+
+        # Exponential backoff with jitter (max 60s)
+        delay = min(60, (2**attempt) + random.uniform(0, 1))
+
+        logger.warning(
+            f"{server.name} LSP server crashed. Restarting in {delay:.2f}s (attempt {attempt + 1})"
+        )
+        await asyncio.sleep(delay)
+
+        # Reset state before restart
+        server.initialized = False
+        server.client = None
+        server.process = None
+        server.pid = None
+
+        try:
+            await self._start_server(server)
+        except Exception as e:
+            logger.error(f"Restart failed for {server.name}: {e}")
+
+    async def _health_check_server(self, server: LSPServer) -> bool:
+        """
+        Perform health check on an LSP server.
+
+        Args:
+            server: LSPServer to check
+
+        Returns:
+            True if server is healthy, False otherwise
+        """
+        if not server.initialized or not server.client:
+            return False
+
+        try:
+            # Simple health check: send initialize request
+            # Most servers respond to repeated initialize calls
+            init_params = InitializeParams(
+                process_id=None, root_uri=None, capabilities=ClientCapabilities()
+            )
+            response = await asyncio.wait_for(
+                server.client.protocol.send_request_async("initialize", init_params),
+                timeout=LSP_CONFIG["health_check_timeout"],
+            )
+            logger.debug(f"{server.name} LSP server health check passed")
+            return True
+        except asyncio.TimeoutError:
+            logger.warning(f"{server.name} LSP server health check timed out")
+            return False
+        except Exception as e:
+            logger.warning(f"{server.name} LSP server health check failed: {e}")
+            return False
+
+    async def _shutdown_single_server(self, name: str, server: LSPServer):
+        """
+        Gracefully shutdown a single LSP server.
+
+        Args:
+            name: Server name (key)
+            server: LSPServer instance
+        """
+        if not server.initialized or not server.client:
+            return
+
+        try:
+            logger.info(f"Shutting down {name} LSP server")
+
+            # LSP protocol shutdown request
+            try:
+                await asyncio.wait_for(
+                    server.client.protocol.send_request_async("shutdown", None), timeout=5.0
+                )
+            except asyncio.TimeoutError:
+                logger.warning(f"{name} LSP server shutdown request timed out")
+
+            # Send exit notification
+            server.client.protocol.notify("exit", None)
+
+            # Stop the client
+            await server.client.stop()
+
+            # Terminate subprocess using stored process reference
+            if server.process is not None:
+                try:
+                    if server.process.returncode is not None:
+                        logger.debug(f"{name} already exited (code {server.process.returncode})")
+                    else:
+                        server.process.terminate()
+                        try:
+                            await asyncio.wait_for(server.process.wait(), timeout=2.0)
+                        except asyncio.TimeoutError:
+                            server.process.kill()
+                            await asyncio.wait_for(server.process.wait(), timeout=1.0)
+                except Exception as e:
+                    logger.warning(f"Error terminating {name}: {e}")
+
+            # Mark as uninitialized
+            server.initialized = False
+            server.client = None
+            server.process = None
+            server.pid = None
+
+        except Exception as e:
+            logger.error(f"Error shutting down {name} LSP server: {e}")
+
+    async def _background_health_monitor(self):
+        """
+        Background task for health checking and idle server shutdown.
+
+        Runs periodically to:
+        - Check health of running servers
+        - Shutdown idle servers
+        - Restart crashed servers
+        """
+        logger.info("LSP health monitor task started")
+        try:
+            while True:
+                await asyncio.sleep(LSP_CONFIG["health_check_interval"])
+
+                current_time = time.time()
+                idle_threshold = current_time - LSP_CONFIG["idle_timeout"]
+
+                async with self._lock:
+                    for name, server in self._servers.items():
+                        if not server.initialized or not server.client:
+                            continue
+
+                        # Check if server is idle
+                        if server.last_used < idle_threshold:
+                            logger.info(
+                                f"{name} LSP server idle for {(current_time - server.last_used) / 60:.1f} minutes, shutting down"
+                            )
+                            await self._shutdown_single_server(name, server)
+                            continue
+
+                        # Health check for active servers
+                        is_healthy = await self._health_check_server(server)
+                        if not is_healthy:
+                            logger.warning(f"{name} LSP server health check failed, restarting")
+                            await self._shutdown_single_server(name, server)
+                            try:
+                                await self._start_server(server)
+                            except Exception as e:
+                                logger.error(f"Failed to restart {name} LSP server: {e}")
+
+        except asyncio.CancelledError:
+            logger.info("LSP health monitor task cancelled")
+            raise
+        except Exception as e:
+            logger.error(f"LSP health monitor task error: {e}", exc_info=True)
+
+    def get_status(self) -> dict:
+        """Get status of managed servers including idle information."""
+        current_time = time.time()
+        status = {}
+        for name, server in self._servers.items():
+            idle_seconds = current_time - server.last_used
+            uptime_seconds = current_time - server.created_at if server.created_at else 0
+            status[name] = {
+                "running": server.initialized and server.client is not None,
+                "pid": server.pid,
+                "command": " ".join(server.command),
+                "restarts": self._restart_attempts.get(name, 0),
+                "idle_seconds": idle_seconds,
+                "idle_minutes": idle_seconds / 60.0,
+                "uptime_seconds": uptime_seconds,
+            }
+        return status
+
+    async def shutdown(self):
+        """
+        Gracefully shutdown all LSP servers.
+
+        Implements:
+        - Health monitor cancellation
+        - LSP protocol shutdown (shutdown request + exit notification)
+        - Pending task cancellation
+        - Process cleanup with timeout
+        """
+        logger.info("Shutting down LSP manager...")
+
+        # Cancel health monitor task
+        if self._health_monitor_task and not self._health_monitor_task.done():
+            logger.info("Cancelling health monitor task")
+            self._health_monitor_task.cancel()
+            try:
+                await self._health_monitor_task
+            except asyncio.CancelledError:
+                pass
+
+        async with self._lock:
+            for name, server in self._servers.items():
+                await self._shutdown_single_server(name, server)
+
+        logger.info("LSP manager shutdown complete")
+
+
+# Singleton accessor
+_manager_instance: Optional[LSPManager] = None
+_manager_lock = threading.Lock()
+
+
+def get_lsp_manager() -> LSPManager:
+    """Get the global LSP manager singleton."""
+    global _manager_instance
+    if _manager_instance is None:
+        with _manager_lock:
+            # Double-check pattern to avoid race condition
+            if _manager_instance is None:
+                _manager_instance = LSPManager()
+    return _manager_instance