llmcode-cli 1.0.0__py3-none-any.whl

llm_code/mcp/manager.py

@@ -0,0 +1,214 @@
+"""McpServerManager: lifecycle management for MCP server connections."""
+from __future__ import annotations
+
+import asyncio
+import warnings
+
+from llm_code.logging import get_logger
+from llm_code.mcp.bridge import McpToolBridge
+from llm_code.mcp.client import McpClient
+from llm_code.mcp.health import MCPHealthChecker
+from llm_code.mcp.transport import HttpTransport, McpTransport, SseTransport, StdioTransport, WebSocketTransport
+from llm_code.mcp.types import McpServerConfig
+from llm_code.tools.registry import ToolRegistry
+
+logger = get_logger(__name__)
+
+# Exponential backoff: 5s, 10s, 20s, 40s → capped at 60s
+_BACKOFF_BASE = 5.0
+_BACKOFF_MAX = 60.0
+
+
+def _backoff_delay(attempt: int) -> float:
+    """Return the delay in seconds for *attempt* (0-indexed), capped at _BACKOFF_MAX."""
+    return min(_BACKOFF_BASE * (2 ** attempt), _BACKOFF_MAX)
+
+
+class McpServerManager:
+    """Manages the lifecycle of MCP server connections and tool registration."""
+
+    def __init__(self) -> None:
+        self._transports: dict[str, McpTransport] = {}
+        self._clients: dict[str, McpClient] = {}
+        self._configs: dict[str, McpServerConfig] = {}
+        self._instructions: dict[str, str] = {}
+        self._health: MCPHealthChecker = MCPHealthChecker()
+        # Track consecutive reconnect failures for backoff
+        self._reconnect_failures: dict[str, int] = {}
+
+    # ------------------------------------------------------------------
+    # Connection management
+    # ------------------------------------------------------------------
+
+    async def start_server(self, name: str, config: McpServerConfig) -> McpClient:
+        """Start a single MCP server, returning an initialised McpClient."""
+        logger.debug("Starting MCP server: %s", name)
+        transport = self._build_transport(config)
+        await transport.start()
+        client = McpClient(transport)
+        info = await client.initialize()
+        self._transports[name] = transport
+        self._clients[name] = client
+        self._configs[name] = config
+        self._reconnect_failures[name] = 0
+
+        # Extract server instructions from capabilities if present
+        capabilities = info.capabilities or {}
+        instructions = capabilities.get("instructions", "")
+        if instructions:
+            self._instructions[name] = instructions
+
+        logger.debug("MCP server started: %s", name)
+        return client
+
+    async def start_all(self, configs: dict[str, McpServerConfig]) -> None:
+        """Start all servers defined in *configs*, logging warnings on failure."""
+        for name, config in configs.items():
+            try:
+                await self.start_server(name, config)
+            except Exception as exc:  # noqa: BLE001
+                logger.warning("Failed to start MCP server '%s': %s", name, exc)
+                warnings.warn(
+                    f"Failed to start MCP server '{name}': {exc}",
+                    stacklevel=2,
+                )
+
+    async def stop_all(self) -> None:
+        """Close all active clients and clear internal state."""
+        logger.debug("Stopping all MCP servers (%d active)", len(self._clients))
+        self._health.stop_monitor()
+        for name, client in list(self._clients.items()):
+            try:
+                await client.close()
+                logger.debug("MCP server stopped: %s", name)
+            except Exception as exc:  # noqa: BLE001
+                logger.warning("Error stopping MCP server '%s': %s", name, exc)
+        self._clients.clear()
+        self._transports.clear()
+
+    def get_client(self, name: str) -> McpClient | None:
+        """Return the client for *name*, or None if not registered."""
+        return self._clients.get(name)
+
+    def get_all_instructions(self) -> dict[str, str]:
+        """Return a mapping of server name → instructions for all servers that provided them."""
+        return {k: v for k, v in self._instructions.items() if v}
+
+    # ------------------------------------------------------------------
+    # Health checking
+    # ------------------------------------------------------------------
+
+    @property
+    def health(self) -> MCPHealthChecker:
+        """Return the :class:`MCPHealthChecker` instance."""
+        return self._health
+
+    async def check_server_health(self, name: str) -> bool:
+        """Check health of a single server.  Returns True if alive."""
+        client = self._clients.get(name)
+        if client is None:
+            return False
+        status = await self._health.check_server(name, client)
+        return status.alive
+
+    async def check_all_health(self):  # type: ignore[return]
+        """Check health of all connected servers concurrently."""
+        return await self._health.check_all(self._clients)
+
+    def start_health_monitor(self, interval: float = 60.0) -> None:
+        """Start background health monitoring for all current clients."""
+        self._health.start_background_monitor(self._clients, interval=interval)
+
+    async def ensure_healthy(self, name: str) -> McpClient:
+        """Return a healthy client for *name*, attempting reconnection once on failure.
+
+        Raises :class:`RuntimeError` if the server is not connected or cannot be
+        reconnected.  Uses exponential backoff on repeated reconnect failures.
+        """
+        client = self._clients.get(name)
+        if client is None:
+            raise RuntimeError(f"MCP server '{name}' is not connected")
+
+        # Quick health probe
+        status = await self._health.check_server(name, client)
+        if status.alive:
+            self._reconnect_failures[name] = 0
+            return client
+
+        # Server appears unhealthy — attempt one reconnect
+        logger.warning("MCP server '%s' is unhealthy (%s), attempting reconnect", name, status.error)
+        config = self._configs.get(name)
+        if config is None:
+            raise RuntimeError(f"No config stored for MCP server '{name}' — cannot reconnect")
+
+        failures = self._reconnect_failures.get(name, 0)
+        if failures > 0:
+            delay = _backoff_delay(failures - 1)
+            logger.debug("Backoff %.0fs before reconnecting '%s' (attempt %d)", delay, name, failures + 1)
+            await asyncio.sleep(delay)
+
+        try:
+            # Close stale connection
+            try:
+                await client.close()
+            except Exception:  # noqa: BLE001
+                pass
+
+            new_client = await self.start_server(name, config)
+            self._reconnect_failures[name] = 0
+            logger.info("MCP server '%s' reconnected successfully", name)
+            return new_client
+        except Exception as exc:  # noqa: BLE001
+            self._reconnect_failures[name] = failures + 1
+            raise RuntimeError(
+                f"Failed to reconnect MCP server '{name}': {exc}"
+            ) from exc
+
+    # ------------------------------------------------------------------
+    # Tool registration
+    # ------------------------------------------------------------------
+
+    async def register_all_tools(self, registry: ToolRegistry) -> int:
+        """Discover and register MCP tools from all active servers.
+
+        Returns the total number of tools registered.
+        """
+        total = 0
+        for server_name, client in self._clients.items():
+            tools = await client.list_tools()
+            for mcp_tool in tools:
+                bridge = McpToolBridge(server_name, mcp_tool, client)
+                registry.register(bridge)
+                total += 1
+        return total
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _build_transport(config: McpServerConfig) -> McpTransport:
+        """Build the appropriate transport from *config*.
+
+        Supported ``transport_type`` values:
+        - ``"stdio"`` (default) — subprocess stdin/stdout
+        - ``"http"`` — HTTP POST requests
+        - ``"sse"`` — Server-Sent Events over HTTP
+        - ``"ws"`` / ``"websocket"`` — WebSocket (requires ``websockets`` package)
+        """
+        if config.transport_type == "http" and config.url:
+            return HttpTransport(url=config.url, headers=config.headers)
+        if config.transport_type == "sse" and config.url:
+            return SseTransport(url=config.url, headers=config.headers)
+        if config.transport_type in ("ws", "websocket") and config.url:
+            return WebSocketTransport(url=config.url, headers=config.headers)
+        if config.command:
+            return StdioTransport(
+                command=config.command,
+                args=config.args,
+                env=config.env,
+            )
+        raise ValueError(
+            f"Cannot build transport from config: transport_type={config.transport_type!r}, "
+            f"command={config.command!r}, url={config.url!r}"
+        )

llm_code/mcp/oauth.py

@@ -0,0 +1,219 @@
+"""OAuth 2.0 Authorization Code + PKCE flow for MCP server authentication."""
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import json
+import secrets
+import threading
+import time
+import urllib.parse
+import webbrowser
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+
+import httpx
+
+
+@dataclass
+class OAuthToken:
+    access_token: str
+    token_type: str = "Bearer"
+    expires_at: float = 0.0
+    refresh_token: str = ""
+    scope: str = ""
+
+    @property
+    def is_expired(self) -> bool:
+        return time.time() >= self.expires_at - 60  # 60s buffer
+
+    def to_dict(self) -> dict:
+        return {
+            "access_token": self.access_token,
+            "token_type": self.token_type,
+            "expires_at": self.expires_at,
+            "refresh_token": self.refresh_token,
+            "scope": self.scope,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> OAuthToken:
+        return cls(**{k: data[k] for k in cls.__dataclass_fields__ if k in data})
+
+
+class OAuthClient:
+    """OAuth 2.0 with PKCE for MCP server authentication."""
+
+    def __init__(
+        self,
+        client_id: str,
+        authorization_url: str,
+        token_url: str,
+        redirect_uri: str = "http://localhost:9876/callback",
+        scope: str = "",
+    ) -> None:
+        self._client_id = client_id
+        self._auth_url = authorization_url
+        self._token_url = token_url
+        self._redirect_uri = redirect_uri
+        self._scope = scope
+        self._token_dir = Path.home() / ".llm-code" / "tokens"
+        self._token_dir.mkdir(parents=True, exist_ok=True)
+
+    def get_token(self, server_name: str) -> OAuthToken | None:
+        """Get a valid token, refreshing if needed."""
+        token = self._load_token(server_name)
+        if token and not token.is_expired:
+            return token
+        if token and token.refresh_token:
+            refreshed = self._refresh_token(token)
+            if refreshed:
+                self._save_token(server_name, refreshed)
+                return refreshed
+        return None
+
+    async def authorize(self, server_name: str) -> OAuthToken:
+        """Run the full OAuth PKCE flow."""
+        # Generate PKCE challenge
+        verifier = secrets.token_urlsafe(64)
+        challenge = (
+            base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest())
+            .rstrip(b"=")
+            .decode()
+        )
+        state = secrets.token_urlsafe(32)
+
+        # Build authorization URL
+        params: dict[str, str] = {
+            "client_id": self._client_id,
+            "response_type": "code",
+            "redirect_uri": self._redirect_uri,
+            "state": state,
+            "code_challenge": challenge,
+            "code_challenge_method": "S256",
+        }
+        if self._scope:
+            params["scope"] = self._scope
+
+        auth_url = f"{self._auth_url}?{urllib.parse.urlencode(params)}"
+
+        # Open browser
+        print("Opening browser for authorization...")
+        print(f"If it doesn't open, visit: {auth_url}")
+        webbrowser.open(auth_url)
+
+        # Start local server to receive callback
+        code = await self._wait_for_callback(state)
+
+        # Exchange code for token
+        token = await self._exchange_code(code, verifier)
+        self._save_token(server_name, token)
+        return token
+
+    async def _exchange_code(self, code: str, verifier: str) -> OAuthToken:
+        """Exchange authorization code for tokens."""
+        async with httpx.AsyncClient() as client:
+            resp = await client.post(
+                self._token_url,
+                data={
+                    "grant_type": "authorization_code",
+                    "client_id": self._client_id,
+                    "code": code,
+                    "redirect_uri": self._redirect_uri,
+                    "code_verifier": verifier,
+                },
+            )
+            resp.raise_for_status()
+            data = resp.json()
+
+        return OAuthToken(
+            access_token=data["access_token"],
+            token_type=data.get("token_type", "Bearer"),
+            expires_at=time.time() + data.get("expires_in", 3600),
+            refresh_token=data.get("refresh_token", ""),
+            scope=data.get("scope", ""),
+        )
+
+    def _refresh_token(self, token: OAuthToken) -> OAuthToken | None:
+        """Refresh an expired token synchronously."""
+        try:
+            with httpx.Client() as client:
+                resp = client.post(
+                    self._token_url,
+                    data={
+                        "grant_type": "refresh_token",
+                        "client_id": self._client_id,
+                        "refresh_token": token.refresh_token,
+                    },
+                )
+                resp.raise_for_status()
+                data = resp.json()
+            return OAuthToken(
+                access_token=data["access_token"],
+                token_type=data.get("token_type", "Bearer"),
+                expires_at=time.time() + data.get("expires_in", 3600),
+                refresh_token=data.get("refresh_token", token.refresh_token),
+                scope=data.get("scope", token.scope),
+            )
+        except Exception:  # noqa: BLE001
+            return None
+
+    async def _wait_for_callback(self, expected_state: str) -> str:
+        """Start a local HTTP server and wait for the OAuth callback."""
+        code_holder: list[str] = []
+
+        class CallbackHandler(BaseHTTPRequestHandler):
+            def do_GET(self) -> None:  # noqa: N802
+                params = urllib.parse.parse_qs(
+                    urllib.parse.urlparse(self.path).query
+                )
+                if (
+                    params.get("state", [""])[0] == expected_state
+                    and "code" in params
+                ):
+                    code_holder.append(params["code"][0])
+                    self.send_response(200)
+                    self.end_headers()
+                    self.wfile.write(
+                        b"Authorization successful! You can close this tab."
+                    )
+                else:
+                    self.send_response(400)
+                    self.end_headers()
+                    self.wfile.write(b"Invalid callback.")
+
+            def log_message(self, *args: object) -> None:  # noqa: ANN002
+                pass  # Suppress log output
+
+        port = int(urllib.parse.urlparse(self._redirect_uri).port or 9876)
+        server = HTTPServer(("localhost", port), CallbackHandler)
+
+        thread = threading.Thread(target=server.handle_request, daemon=True)
+        thread.start()
+
+        # Wait for callback (max 120 seconds)
+        for _ in range(1200):
+            if code_holder:
+                break
+            await asyncio.sleep(0.1)
+
+        server.server_close()
+
+        if not code_holder:
+            raise TimeoutError("OAuth callback timed out")
+        return code_holder[0]
+
+    def _save_token(self, server_name: str, token: OAuthToken) -> None:
+        path = self._token_dir / f"{server_name}.json"
+        path.write_text(json.dumps(token.to_dict()))
+
+    def _load_token(self, server_name: str) -> OAuthToken | None:
+        path = self._token_dir / f"{server_name}.json"
+        if path.exists():
+            try:
+                return OAuthToken.from_dict(json.loads(path.read_text()))
+            except Exception:  # noqa: BLE001
+                return None
+        return None

llm_code/mcp/transport.py

@@ -0,0 +1,254 @@
+"""MCP transport layer: ABC, StdioTransport, HttpTransport, SseTransport, WebSocketTransport."""
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from abc import ABC, abstractmethod
+from typing import Any
+
+import httpx
+
+
+class McpTransport(ABC):
+    """Abstract base class for MCP transports."""
+
+    @abstractmethod
+    async def start(self) -> None:
+        """Start the transport (connect or launch subprocess)."""
+
+    @abstractmethod
+    async def send(self, message: dict[str, Any]) -> None:
+        """Send a JSON message."""
+
+    @abstractmethod
+    async def receive(self) -> dict[str, Any]:
+        """Receive and return the next JSON message."""
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Close and clean up the transport."""
+
+
+class StdioTransport(McpTransport):
+    """MCP transport that communicates via subprocess stdin/stdout."""
+
+    RECEIVE_TIMEOUT = 30.0
+    CLOSE_WAIT_TIMEOUT = 5.0
+
+    def __init__(
+        self,
+        command: str,
+        args: tuple[str, ...] = (),
+        env: dict[str, str] | None = None,
+    ) -> None:
+        self.command = command
+        self.args = args
+        self.env = env
+        self._process: asyncio.subprocess.Process | None = None
+
+    async def start(self) -> None:
+        """Launch the subprocess with merged environment."""
+        merged_env = {**os.environ}
+        if self.env:
+            merged_env.update(self.env)
+
+        self._process = await asyncio.create_subprocess_exec(
+            self.command,
+            *self.args,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=merged_env,
+        )
+
+    async def send(self, message: dict[str, Any]) -> None:
+        """Write JSON + newline to subprocess stdin."""
+        if self._process is None or self._process.stdin is None:
+            raise RuntimeError("Transport not started")
+        data = json.dumps(message) + "\n"
+        self._process.stdin.write(data.encode())
+        await self._process.stdin.drain()
+
+    async def receive(self) -> dict[str, Any]:
+        """Read a line from subprocess stdout with timeout, parse JSON."""
+        if self._process is None or self._process.stdout is None:
+            raise RuntimeError("Transport not started")
+        line = await asyncio.wait_for(
+            self._process.stdout.readline(),
+            timeout=self.RECEIVE_TIMEOUT,
+        )
+        return json.loads(line.decode().strip())
+
+    async def close(self) -> None:
+        """Terminate subprocess gracefully, kill if needed."""
+        if self._process is None:
+            return
+        process = self._process
+        self._process = None
+
+        try:
+            if process.stdin and not process.stdin.is_closing():
+                process.stdin.close()
+        except Exception:
+            pass
+
+        try:
+            process.terminate()
+        except (ProcessLookupError, OSError):
+            pass
+
+        try:
+            await asyncio.wait_for(process.wait(), timeout=self.CLOSE_WAIT_TIMEOUT)
+        except asyncio.TimeoutError:
+            try:
+                process.kill()
+            except (ProcessLookupError, OSError):
+                pass
+            try:
+                await process.wait()
+            except Exception:
+                pass
+
+
+class HttpTransport(McpTransport):
+    """MCP transport that communicates via HTTP POST requests."""
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        self.url = url
+        self.headers = headers
+        self._client: httpx.AsyncClient | None = None
+        self._last_response: dict[str, Any] | None = None
+
+    async def start(self) -> None:
+        """Create the HTTP client."""
+        self._client = httpx.AsyncClient(headers=self.headers or {})
+
+    async def send(self, message: dict[str, Any]) -> None:
+        """POST the JSON message to the server URL and store the response."""
+        if self._client is None:
+            raise RuntimeError("Transport not started")
+        response = await self._client.post(self.url, json=message)
+        response.raise_for_status()
+        self._last_response = response.json()
+
+    async def receive(self) -> dict[str, Any]:
+        """Return the stored response from the last send()."""
+        if self._last_response is None:
+            raise RuntimeError("No response available; call send() first")
+        result = self._last_response
+        self._last_response = None
+        return result
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        if self._client is None:
+            return
+        client = self._client
+        self._client = None
+        await client.aclose()
+
+
+class SseTransport(McpTransport):
+    """SSE (Server-Sent Events) transport for MCP.
+
+    Server-Sent Events are server→client by nature; outbound messages are
+    sent via HTTP POST to the same URL, and responses are queued for
+    retrieval via ``receive()``.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        self._url = url
+        self._headers: dict[str, str] = headers or {}
+        self._client: httpx.AsyncClient | None = None
+        self._response_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
+
+    async def start(self) -> None:
+        """Create the HTTP client with SSE accept header."""
+        self._client = httpx.AsyncClient(
+            headers={"Accept": "text/event-stream", **self._headers},
+            timeout=httpx.Timeout(30.0),
+        )
+
+    async def send(self, message: dict[str, Any]) -> None:
+        """Send a message via HTTP POST and enqueue the response."""
+        if self._client is None:
+            raise RuntimeError("Transport not started")
+        resp = await self._client.post(self._url, json=message)
+        resp.raise_for_status()
+        await self._response_queue.put(resp.json())
+
+    async def receive(self) -> dict[str, Any]:
+        """Return the next queued response."""
+        return await self._response_queue.get()
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        if self._client is None:
+            return
+        client = self._client
+        self._client = None
+        await client.aclose()
+
+
+class WebSocketTransport(McpTransport):
+    """WebSocket transport for MCP.
+
+    Requires the optional ``websockets`` package::
+
+        pip install websockets
+        # or: pip install llm-code[websocket]
+    """
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        self._url = url
+        self._headers: dict[str, str] = headers or {}
+        self._ws: Any = None  # websockets.WebSocketClientProtocol
+
+    async def start(self) -> None:
+        """Connect to the WebSocket server."""
+        try:
+            import websockets  # type: ignore[import-untyped]
+        except ImportError as exc:
+            raise ImportError(
+                "WebSocket transport requires the 'websockets' package: "
+                "pip install websockets"
+            ) from exc
+
+        self._ws = await websockets.connect(
+            self._url,
+            additional_headers=self._headers,
+        )
+
+    async def send(self, message: dict[str, Any]) -> None:
+        """Send a JSON message over the WebSocket."""
+        if self._ws is None:
+            raise RuntimeError("WebSocket not connected")
+        await self._ws.send(json.dumps(message))
+
+    async def receive(self) -> dict[str, Any]:
+        """Receive and parse the next JSON message from the WebSocket."""
+        if self._ws is None:
+            raise RuntimeError("WebSocket not connected")
+        data = await self._ws.recv()
+        return json.loads(data)  # type: ignore[no-any-return]
+
+    async def close(self) -> None:
+        """Close the WebSocket connection."""
+        if self._ws is None:
+            return
+        ws = self._ws
+        self._ws = None
+        await ws.close()