server_manager.py

from __future__ import annotations

import os
import subprocess
import threading
from dataclasses import dataclass
from typing import Dict, Optional

from core.mcp.client import StdioMCPClient
from core.mcp.config import MCPServerConfig, load_mcp_servers


@dataclass
class _ServerHandle:
    cfg: MCPServerConfig
    proc: subprocess.Popen
    client: StdioMCPClient


class MCPServerManager:
    """Starts and manages stdio MCP server subprocesses (demo: one process per server id)."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._servers_by_id: Dict[str, MCPServerConfig] = {s.id: s for s in load_mcp_servers()}
        self._handles: Dict[str, _ServerHandle] = {}

    def list_server_ids(self) -> list[str]:
        with self._lock:
            return sorted(self._servers_by_id.keys())

    def get_client(self, *, server_id: str) -> StdioMCPClient:
        with self._lock:
            if server_id in self._handles:
                return self._handles[server_id].client
            cfg = self._servers_by_id[server_id]
            if cfg.transport != "stdio":
                raise ValueError(f"Unsupported MCP transport: {cfg.transport}")

            env = dict(**(cfg.env or {}))
            proc = subprocess.Popen(
                cfg.command,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                cwd=cfg.cwd,
                env={**os.environ, **env},
            )
            assert proc.stdin is not None
            assert proc.stdout is not None
            client = StdioMCPClient(server_id=server_id, stdin=proc.stdin, stdout=proc.stdout)
            # initialize eagerly to fail fast
            client.initialize()
            self._handles[server_id] = _ServerHandle(cfg=cfg, proc=proc, client=client)
            return client

    def shutdown(self) -> None:
        with self._lock:
            for hid, h in list(self._handles.items()):
                try:
                    h.proc.terminate()
                except Exception:
                    pass
                self._handles.pop(hid, None)