pylance-mcp-server 1.0.0

package/mcp_server/pylance_bridge.py

@@ -0,0 +1,579 @@
+"""
+Pylance Bridge Module
+
+Spawns and manages communication with the Pyright language server via subprocess.
+Implements LSP JSON-RPC protocol over stdio.
+"""
+
+import asyncio
+import json
+import logging
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+import threading
+import queue
+
+from opentelemetry import trace
+
+logger = logging.getLogger(__name__)
+tracer = trace.get_tracer(__name__)
+
+
+class PylanceBridge:
+    """Manages a Pyright language server subprocess and handles LSP communication."""
+
+    def __init__(self, workspace_root: str):
+        """
+        Initialize the Pylance bridge.
+
+        Args:
+            workspace_root: Absolute path to the workspace root directory
+        """
+        self.workspace_root = Path(workspace_root).resolve()
+        self.process: Optional[subprocess.Popen] = None
+        self.message_id = 0
+        self.pending_responses: Dict[int, queue.Queue] = {}
+        self.reader_thread: Optional[threading.Thread] = None
+        self.running = False
+        self.initialized = False
+        self._lock = threading.Lock()
+        
+        # Detect Python virtual environment
+        self.python_path = self._detect_python_venv()
+        if self.python_path:
+            logger.info(f"Detected Python environment: {self.python_path}")
+        else:
+            logger.warning("No virtual environment detected, using system Python")
+
+    def _detect_python_venv(self) -> Optional[str]:
+        """
+        Detect and return the Python executable path from virtual environment.
+        If no venv exists, creates one automatically.
+        
+        Checks in order:
+        1. .venv/Scripts/python.exe (Windows) or .venv/bin/python (Unix)
+        2. venv/Scripts/python.exe (Windows) or venv/bin/python (Unix)
+        3. env/Scripts/python.exe (Windows) or env/bin/python (Unix)
+        4. .env/Scripts/python.exe (Windows) or .env/bin/python (Unix)
+        5. If none found, creates .venv automatically
+        6. System Python (fallback if creation fails)
+        
+        Returns:
+            Path to Python executable or None if not found
+        """
+        venv_names = ['.venv', 'venv', 'env', '.env']
+        
+        if sys.platform == 'win32':
+            python_names = ['python.exe', 'python3.exe']
+            venv_subdir = 'Scripts'
+        else:
+            python_names = ['python3', 'python']
+            venv_subdir = 'bin'
+        
+        # Check for existing virtual environments in workspace
+        for venv_name in venv_names:
+            venv_path = self.workspace_root / venv_name
+            if venv_path.exists() and venv_path.is_dir():
+                for python_name in python_names:
+                    python_exe = venv_path / venv_subdir / python_name
+                    if python_exe.exists():
+                        logger.info(f"Found existing virtual environment: {venv_path}")
+                        return str(python_exe.resolve())
+        
+        # No venv found - create one automatically
+        logger.info("No virtual environment found. Creating .venv automatically...")
+        created_python = self._create_venv()
+        if created_python:
+            return created_python
+        
+        # Fallback to system Python
+        for python_name in python_names:
+            try:
+                result = subprocess.run(
+                    ['which' if sys.platform != 'win32' else 'where', python_name],
+                    capture_output=True,
+                    text=True,
+                    timeout=5
+                )
+                if result.returncode == 0 and result.stdout.strip():
+                    python_path = result.stdout.strip().split('\n')[0]
+                    logger.info(f"Using system Python: {python_path}")
+                    return python_path
+            except Exception as e:
+                logger.debug(f"Error finding {python_name}: {e}")
+        
+        return None
+
+    def _create_venv(self) -> Optional[str]:
+        """
+        Create a new virtual environment in .venv directory.
+        
+        Returns:
+            Path to created Python executable, or None if creation failed
+        """
+        venv_path = self.workspace_root / '.venv'
+        
+        # Find system Python to create venv
+        if sys.platform == 'win32':
+            python_names = ['python', 'python3', 'py']
+        else:
+            python_names = ['python3', 'python']
+        
+        system_python = None
+        for python_name in python_names:
+            try:
+                result = subprocess.run(
+                    ['which' if sys.platform != 'win32' else 'where', python_name],
+                    capture_output=True,
+                    text=True,
+                    timeout=5
+                )
+                if result.returncode == 0 and result.stdout.strip():
+                    system_python = result.stdout.strip().split('\n')[0]
+                    break
+            except Exception as e:
+                logger.debug(f"Error finding {python_name}: {e}")
+        
+        if not system_python:
+            logger.error("Cannot create venv: No Python executable found")
+            return None
+        
+        logger.info(f"Creating virtual environment at {venv_path} using {system_python}")
+        
+        try:
+            # Create venv using: python -m venv .venv
+            result = subprocess.run(
+                [system_python, '-m', 'venv', str(venv_path)],
+                capture_output=True,
+                text=True,
+                timeout=60,
+                cwd=str(self.workspace_root)
+            )
+            
+            if result.returncode != 0:
+                logger.error(f"Failed to create venv: {result.stderr}")
+                return None
+            
+            # Verify venv was created
+            if sys.platform == 'win32':
+                python_exe = venv_path / 'Scripts' / 'python.exe'
+            else:
+                python_exe = venv_path / 'bin' / 'python3'
+                if not python_exe.exists():
+                    python_exe = venv_path / 'bin' / 'python'
+            
+            if python_exe.exists():
+                logger.info(f"✓ Successfully created virtual environment: {venv_path}")
+                
+                # Install common packages if requirements.txt exists
+                requirements_file = self.workspace_root / 'requirements.txt'
+                if requirements_file.exists():
+                    logger.info("Found requirements.txt, installing dependencies...")
+                    self._install_requirements(str(python_exe), str(requirements_file))
+                else:
+                    logger.info("No requirements.txt found. Virtual environment created with no packages.")
+                
+                return str(python_exe.resolve())
+            else:
+                logger.error(f"Venv created but Python executable not found at {python_exe}")
+                return None
+                
+        except subprocess.TimeoutExpired:
+            logger.error("Timeout creating virtual environment (>60s)")
+            return None
+        except Exception as e:
+            logger.error(f"Error creating virtual environment: {e}")
+            return None
+
+    def _install_requirements(self, python_exe: str, requirements_file: str) -> None:
+        """
+        Install packages from requirements.txt into the virtual environment.
+        
+        Args:
+            python_exe: Path to venv Python executable
+            requirements_file: Path to requirements.txt
+        """
+        try:
+            logger.info(f"Installing packages from {requirements_file}...")
+            
+            # Use pip from the venv: python -m pip install -r requirements.txt
+            result = subprocess.run(
+                [python_exe, '-m', 'pip', 'install', '-r', requirements_file],
+                capture_output=True,
+                text=True,
+                timeout=300,  # 5 minutes timeout for package installation
+                cwd=str(self.workspace_root)
+            )
+            
+            if result.returncode == 0:
+                logger.info("✓ Successfully installed dependencies from requirements.txt")
+            else:
+                logger.warning(f"Failed to install some dependencies: {result.stderr}")
+                
+        except subprocess.TimeoutExpired:
+            logger.warning("Package installation timed out (>5 minutes). Some packages may not be installed.")
+        except Exception as e:
+            logger.warning(f"Error installing requirements: {e}")
+
+
+    def _create_pyrightconfig(self) -> None:
+        """
+        Create or update pyrightconfig.json to use detected Python environment.
+        """
+        config_path = self.workspace_root / 'pyrightconfig.json'
+        
+        config = {}
+        if config_path.exists():
+            try:
+                with open(config_path, 'r', encoding='utf-8') as f:
+                    config = json.load(f)
+            except Exception as e:
+                logger.warning(f"Failed to read existing pyrightconfig.json: {e}")
+        
+        # Update Python path if detected
+        if self.python_path:
+            config['pythonPath'] = self.python_path
+        
+        # Set other useful defaults
+        if 'include' not in config:
+            config['include'] = ['.']
+        
+        if 'exclude' not in config:
+            config['exclude'] = [
+                '**/node_modules',
+                '**/__pycache__',
+                '.git',
+                '.venv',
+                'venv',
+                'env',
+                '.env'
+            ]
+        
+        # Write back to file
+        try:
+            with open(config_path, 'w', encoding='utf-8') as f:
+                json.dump(config, f, indent=2)
+            logger.info(f"Created/updated pyrightconfig.json with Python path: {self.python_path}")
+        except Exception as e:
+            logger.error(f"Failed to write pyrightconfig.json: {e}")
+
+    def start(self) -> None:
+        """Start the Pyright language server subprocess."""
+        if self.process and self.process.poll() is None:
+            logger.info("Pyright server already running")
+            return
+
+        # Create/update pyrightconfig.json with detected Python path
+        self._create_pyrightconfig()
+
+        # Find npx or pyright executable
+        npx_cmd = self._find_npx()
+        if not npx_cmd:
+            raise RuntimeError(
+                "npx not found. Please install Node.js: https://nodejs.org/"
+            )
+
+        # Start the language server
+        cmd = [npx_cmd, "pyright-langserver", "--stdio"]
+        logger.info(f"Starting Pyright: {' '.join(cmd)}")
+
+        try:
+            self.process = subprocess.Popen(
+                cmd,
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                cwd=str(self.workspace_root),
+                text=False,  # Use binary mode for proper encoding
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to start Pyright: {e}")
+
+        self.running = True
+        self.reader_thread = threading.Thread(target=self._read_messages, daemon=True)
+        self.reader_thread.start()
+
+        # Initialize the language server
+        self._initialize()
+        logger.info("Pyright language server started successfully")
+
+    def stop(self) -> None:
+        """Stop the Pyright language server subprocess."""
+        self.running = False
+        if self.process:
+            try:
+                self.send_request("shutdown", {})
+                self.send_notification("exit", {})
+                self.process.wait(timeout=5)
+            except Exception as e:
+                logger.warning(f"Error during shutdown: {e}")
+                self.process.kill()
+            finally:
+                self.process = None
+                self.initialized = False
+        logger.info("Pyright language server stopped")
+
+    def _find_npx(self) -> Optional[str]:
+        """Find the npx executable."""
+        # Check if npx is in PATH
+        if sys.platform == "win32":
+            npx_names = ["npx.cmd", "npx.exe", "npx"]
+        else:
+            npx_names = ["npx"]
+
+        for npx_name in npx_names:
+            try:
+                result = subprocess.run(
+                    ["where" if sys.platform == "win32" else "which", npx_name],
+                    capture_output=True,
+                    text=True,
+                )
+                if result.returncode == 0:
+                    return npx_name
+            except Exception:
+                pass
+
+        # Check common locations
+        if sys.platform == "win32":
+            common_paths = [
+                os.path.expandvars(r"%ProgramFiles%\nodejs"),
+                os.path.expandvars(r"%ProgramFiles(x86)%\nodejs"),
+                os.path.expandvars(r"%APPDATA%\npm"),
+            ]
+        else:
+            common_paths = [
+                "/usr/local/bin",
+                "/usr/bin",
+                os.path.expanduser("~/.npm-global/bin"),
+            ]
+
+        for path in common_paths:
+            for npx_name in npx_names:
+                npx_path = os.path.join(path, npx_name)
+                if os.path.exists(npx_path):
+                    return npx_path
+
+        return None
+
+    def _initialize(self) -> None:
+        """Send the initialize request to the language server."""
+        init_params = {
+            "processId": os.getpid(),
+            "rootUri": self.workspace_root.as_uri(),
+            "capabilities": {
+                "textDocument": {
+                    "completion": {"completionItem": {"snippetSupport": True}},
+                    "hover": {"contentFormat": ["markdown", "plaintext"]},
+                    "definition": {"linkSupport": True},
+                    "references": {},
+                    "rename": {"prepareSupport": True},
+                    "diagnostic": {},
+                    "formatting": {},
+                },
+                "workspace": {
+                    "applyEdit": True,
+                    "workspaceEdit": {"documentChanges": True},
+                },
+            },
+            "workspaceFolders": [
+                {"uri": self.workspace_root.as_uri(), "name": self.workspace_root.name}
+            ],
+        }
+        
+        # Add Python path to initialization options if detected
+        if self.python_path:
+            init_params["initializationOptions"] = {
+                "pythonPath": self.python_path,
+                "python": {
+                    "pythonPath": self.python_path
+                }
+            }
+
+        response = self.send_request("initialize", init_params)
+        if "error" in response:
+            raise RuntimeError(f"Failed to initialize: {response['error']}")
+
+        self.send_notification("initialized", {})
+        
+        # Send workspace configuration with Python path
+        if self.python_path:
+            self.send_notification("workspace/didChangeConfiguration", {
+                "settings": {
+                    "python": {
+                        "pythonPath": self.python_path
+                    }
+                }
+            })
+        
+        self.initialized = True
+        logger.info("Pyright language server initialized")
+
+    def _read_messages(self) -> None:
+        """Read messages from the language server stdout."""
+        buffer = b""
+        while self.running and self.process:
+            try:
+                chunk = self.process.stdout.read(1024)
+                if not chunk:
+                    break
+
+                buffer += chunk
+
+                while b"\r\n\r\n" in buffer:
+                    header_end = buffer.index(b"\r\n\r\n")
+                    header = buffer[:header_end].decode("utf-8")
+                    buffer = buffer[header_end + 4 :]
+
+                    # Parse Content-Length
+                    content_length = 0
+                    for line in header.split("\r\n"):
+                        if line.startswith("Content-Length:"):
+                            content_length = int(line.split(":")[1].strip())
+                            break
+
+                    # Read the message body
+                    while len(buffer) < content_length:
+                        chunk = self.process.stdout.read(content_length - len(buffer))
+                        if not chunk:
+                            break
+                        buffer += chunk
+
+                    message_body = buffer[:content_length]
+                    buffer = buffer[content_length:]
+
+                    try:
+                        message = json.loads(message_body.decode("utf-8"))
+                        self._handle_message(message)
+                    except json.JSONDecodeError as e:
+                        logger.error(f"Failed to decode message: {e}")
+
+            except Exception as e:
+                logger.error(f"Error reading messages: {e}")
+                break
+
+    def _handle_message(self, message: Dict[str, Any]) -> None:
+        """Handle incoming message from the language server."""
+        if "id" in message:
+            # This is a response to a request
+            msg_id = message["id"]
+            if msg_id in self.pending_responses:
+                self.pending_responses[msg_id].put(message)
+        else:
+            # This is a notification - log it
+            logger.debug(f"Received notification: {message.get('method')}")
+
+    def send_notification(self, method: str, params: Dict[str, Any]) -> None:
+        """Send a notification to the language server."""
+        message = {"jsonrpc": "2.0", "method": method, "params": params}
+        self._send_message(message)
+
+    def send_request(self, method: str, params: Dict[str, Any], timeout: int = 30) -> Dict[str, Any]:
+        """
+        Send a request to the language server and wait for response.
+
+        Args:
+            method: LSP method name
+            params: Request parameters
+            timeout: Timeout in seconds
+
+        Returns:
+            Response dictionary
+        """
+        with tracer.start_as_current_span(f"lsp.{method}") as span:
+            span.set_attribute("lsp.method", method)
+            span.set_attribute("lsp.timeout", timeout)
+            
+            with self._lock:
+                self.message_id += 1
+                msg_id = self.message_id
+
+            response_queue: queue.Queue = queue.Queue()
+            self.pending_responses[msg_id] = response_queue
+
+            message = {"jsonrpc": "2.0", "id": msg_id, "method": method, "params": params}
+            self._send_message(message)
+
+            try:
+                response = response_queue.get(timeout=timeout)
+                span.set_attribute("lsp.success", True)
+                return response
+            except queue.Empty:
+                span.set_attribute("lsp.success", False)
+                span.set_attribute("lsp.error", "timeout")
+                raise TimeoutError(f"Request {method} timed out after {timeout}s")
+            finally:
+                self.pending_responses.pop(msg_id, None)
+
+    def _send_message(self, message: Dict[str, Any]) -> None:
+        """Send a message to the language server."""
+        if not self.process or not self.process.stdin:
+            raise RuntimeError("Language server not running")
+
+        content = json.dumps(message).encode("utf-8")
+        header = f"Content-Length: {len(content)}\r\n\r\n".encode("utf-8")
+
+        try:
+            self.process.stdin.write(header + content)
+            self.process.stdin.flush()
+        except Exception as e:
+            raise RuntimeError(f"Failed to send message: {e}")
+
+    def _file_uri(self, file_path: str) -> str:
+        """Convert file path to URI."""
+        path = Path(file_path)
+        if not path.is_absolute():
+            path = (self.workspace_root / path).resolve()
+        return path.as_uri()
+
+    def _validate_path(self, file_path: str) -> Path:
+        """Validate that file path is within workspace root."""
+        path = Path(file_path)
+        if not path.is_absolute():
+            path = (self.workspace_root / path).resolve()
+        else:
+            path = path.resolve()
+
+        try:
+            path.relative_to(self.workspace_root)
+        except ValueError:
+            raise ValueError(f"File path {file_path} is outside workspace root")
+
+        return path
+
+    def open_document(self, file_path: str, content: str) -> None:
+        """Open a document in the language server."""
+        with tracer.start_as_current_span("lsp.open_document") as span:
+            span.set_attribute("file_path", file_path)
+            span.set_attribute("content_length", len(content))
+            
+            uri = self._file_uri(file_path)
+            params = {
+                "textDocument": {
+                    "uri": uri,
+                    "languageId": "python",
+                    "version": 1,
+                    "text": content,
+                }
+            }
+            self.send_notification("textDocument/didOpen", params)
+
+    def close_document(self, file_path: str) -> None:
+        """Close a document in the language server."""
+        with tracer.start_as_current_span("lsp.close_document") as span:
+            span.set_attribute("file_path", file_path)
+            
+            uri = self._file_uri(file_path)
+            params = {"textDocument": {"uri": uri}}
+            self.send_notification("textDocument/didClose", params)
+
+    def __enter__(self) -> "PylanceBridge":
+        """Context manager entry."""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Context manager exit."""
+        self.stop()