Agentic-OS/backend/app/mcp_manager/manager.py

"""Clone MCP repos from Gitea and run them as stdio MCP subprocess sidecars."""
from __future__ import annotations

import asyncio
import json
import logging
import os
import shutil
from contextlib import AsyncExitStack
from dataclasses import dataclass, field

from app.config import settings

logger = logging.getLogger(__name__)

MCP_STATUS_KEY = "agentic:mcp:status"


@dataclass
class ServerState:
    name: str
    status: str = "stopped"
    tools: list[dict] = field(default_factory=list)
    error: str | None = None
    session: object | None = None
    _stack: AsyncExitStack | None = field(default=None, repr=False)


async def ensure_git_remote(repo: str, target: str) -> None:
    """Point origin at GITEA_BASE_URL — fixes repos cloned when remote was public HTTPS."""
    if not os.path.isdir(os.path.join(target, ".git")):
        return
    url = settings.gitea_clone_url(repo)
    code, out = await McpManager._run(["git", "remote", "set-url", "origin", url], cwd=target)
    if code != 0:
        raise RuntimeError(f"git remote set-url failed: {out[:500]}")


class McpManager:
    def __init__(self) -> None:
        self.servers: dict[str, ServerState] = {}
        self._lock = asyncio.Lock()
        self._started = False

    async def startup(self) -> None:
        async with self._lock:
            if self._started:
                return
            self._started = True
            os.makedirs(self.repos_dir, exist_ok=True)
            for repo in settings.mcp_repo_list:
                state = ServerState(name=repo)
                self.servers[repo] = state
                try:
                    await self._prepare_repo(repo, state)
                    await self._start_server(repo, state)
                except Exception as exc:  # noqa: BLE001
                    state.status = "error"
                    state.error = str(exc)
                    logger.warning("MCP %s failed to start: %s", repo, exc)
            await self.publish_status()

    async def shutdown(self) -> None:
        async with self._lock:
            for state in self.servers.values():
                await self._stop_server(state)

    @property
    def repos_dir(self) -> str:
        return os.path.join(settings.runtime_dir, "mcp-repos")

    @property
    def local_dir(self) -> str:
        return os.environ.get("MCP_LOCAL_DIR", "/app/mcp-servers")

    def _is_local_repo(self, repo: str) -> bool:
        return os.path.isdir(os.path.join(self.local_dir, repo))

    def _repo_dir(self, repo: str) -> str:
        local_src = os.path.join(self.local_dir, repo)
        if os.path.isdir(local_src):
            dst = os.path.join(settings.runtime_dir, "mcp-local", repo)
            if not os.path.isdir(dst):
                os.makedirs(os.path.dirname(dst), exist_ok=True)
                shutil.copytree(local_src, dst)
            return dst
        return os.path.join(self.repos_dir, repo)

    async def _sync_local_repo(self, repo: str) -> None:
        local_src = os.path.join(self.local_dir, repo)
        if not os.path.isdir(local_src):
            return
        dst = os.path.join(settings.runtime_dir, "mcp-local", repo)
        if os.path.isdir(dst):
            shutil.rmtree(dst)
        shutil.copytree(local_src, dst)

    async def _ensure_git_remote(self, repo: str, target: str) -> None:
        await ensure_git_remote(repo, target)

    async def _prepare_repo(self, repo: str, state: ServerState) -> None:
        if self._is_local_repo(repo):
            await self._sync_local_repo(repo)
            return
        state.status = "cloning"
        target = os.path.join(self.repos_dir, repo)
        url = settings.gitea_clone_url(repo)
        if os.path.isdir(os.path.join(target, ".git")):
            await self._ensure_git_remote(repo, target)
            code, out = await self._run(["git", "pull", "--rebase", "--autostash"], cwd=target)
            if code != 0:
                raise RuntimeError(f"git pull failed: {out[:500]}")
        else:
            code, out = await self._run(["git", "clone", "--depth", "1", url, target])
            if code != 0:
                raise RuntimeError(f"clone failed: {out[:500]}")

    def _detect_start_command(self, repo: str) -> list[str]:
        override = os.environ.get(f"MCP_START_{repo.upper().replace('-', '_')}")
        if override:
            return override.split()
        target = self._repo_dir(repo)
        has_pyproject = os.path.isfile(os.path.join(target, "pyproject.toml"))
        for candidate in ("server.py", "main.py", "app.py", "mcp_server.py", "__main__.py"):
            if os.path.isfile(os.path.join(target, candidate)):
                if has_pyproject:
                    return ["uv", "run", "--directory", target, "python", candidate]
                return ["uv", "run", "--no-project", "--directory", target, "python", candidate]
        module = repo.replace("-", "_")
        if os.path.isdir(os.path.join(target, "src", module)) or os.path.isdir(
            os.path.join(target, module)
        ):
            return ["uv", "run", "--directory", target, "python", "-m", module]
        return ["uv", "run", "--directory", target, repo]

    async def _stop_server(self, state: ServerState) -> None:
        if state._stack is not None:
            try:
                await state._stack.aclose()
            except Exception as exc:  # noqa: BLE001
                logger.debug("stop %s: %s", state.name, exc)
            state._stack = None
        state.session = None
        state.tools = []
        if state.status not in ("error", "cloning"):
            state.status = "stopped"

    def _mcp_subprocess_env(self, repo: str) -> dict[str, str]:
        """Env passed to MCP stdio subprocesses (repo-specific overrides)."""
        env = {**os.environ}
        if repo == "proxmox-mcp":
            verify = os.environ.get("PROXMOX_VERIFY_SSL")
            if verify is None:
                verify = os.environ.get("DEVICE_PROXMOX_VERIFY_SSL", "false")
            env["PROXMOX_VERIFY_SSL"] = verify
        return env

    async def _start_server(self, repo: str, state: ServerState) -> None:
        from mcp import ClientSession, StdioServerParameters
        from mcp.client.stdio import stdio_client

        state.status = "starting"
        state.error = None
        target = self._repo_dir(repo)
        cmd = self._detect_start_command(repo)
        params = StdioServerParameters(
            command=cmd[0],
            args=cmd[1:],
            cwd=target,
            env=self._mcp_subprocess_env(repo),
        )
        stack = AsyncExitStack()
        read, write = await stack.enter_async_context(stdio_client(params))
        session = await stack.enter_async_context(ClientSession(read, write))
        await asyncio.wait_for(session.initialize(), timeout=90)
        tools_resp = await session.list_tools()
        state._stack = stack
        state.session = session
        state.tools = [{"name": t.name, "description": t.description} for t in tools_resp.tools]
        state.status = "loaded"
        logger.info("MCP %s loaded with %d tools", repo, len(state.tools))

    async def reload_server(self, repo: str) -> None:
        async with self._lock:
            state = self.servers.get(repo)
            if not state:
                raise ValueError(f"Unknown MCP server: {repo}")
            await self._stop_server(state)
            await self._start_server(repo, state)
            await self.publish_status()

    async def upgrade_server(self, repo: str) -> None:
        async with self._lock:
            state = self.servers.get(repo)
            if not state:
                raise ValueError(f"Unknown MCP server: {repo}")
            await self._stop_server(state)
            await self._prepare_repo(repo, state)
            await self._start_server(repo, state)
            await self.publish_status()

    async def reload_all(self) -> None:
        async with self._lock:
            for repo in settings.mcp_repo_list:
                state = self.servers.get(repo)
                if not state:
                    continue
                try:
                    await self._stop_server(state)
                    await self._start_server(repo, state)
                except Exception as exc:  # noqa: BLE001
                    state.status = "error"
                    state.error = str(exc)
            await self.publish_status()

    async def process_command_queue_once(self) -> bool:
        from app.mcp_manager.commands import (
            MCP_COMMAND_QUEUE,
            process_mcp_command,
            run_develop_command_background,
            store_mcp_command_result,
        )
        from app.services.events import get_redis

        item = await get_redis().lpop(MCP_COMMAND_QUEUE)
        if not item:
            return False
        cmd = json.loads(item)
        if cmd.get("action") == "develop":
            asyncio.create_task(run_develop_command_background(cmd, self))
            return True
        result = await process_mcp_command(cmd)
        await store_mcp_command_result(cmd["request_id"], result)
        return True

    def find_server_for_tool(self, tool_name: str) -> str | None:
        for name, state in self.servers.items():
            if any(t["name"] == tool_name for t in state.tools):
                return name
        return None

    async def call_tool(self, server: str, tool: str, args: dict) -> dict:
        state = self.servers.get(server)
        if not state or state.status != "loaded" or state.session is None:
            raise RuntimeError(f"MCP server '{server}' is not available")
        result = await state.session.call_tool(tool, args)  # type: ignore[attr-defined]
        content = []
        for block in getattr(result, "content", []) or []:
            text = getattr(block, "text", None)
            if text is not None:
                content.append(text)
        return {
            "isError": getattr(result, "isError", False),
            "text": "\n".join(content),
            "structured": getattr(result, "structuredContent", None),
        }

    def status_snapshot(self) -> list[dict]:
        return [
            {
                "name": s.name,
                "status": s.status,
                "tools": s.tools,
                "error": s.error,
                "source": "local" if self._is_local_repo(s.name) else "git",
            }
            for s in self.servers.values()
        ]

    async def publish_status(self) -> None:
        from app.services.events import get_redis

        await get_redis().set(MCP_STATUS_KEY, json.dumps(self.status_snapshot()))

    @staticmethod
    async def _run(cmd: list[str], cwd: str | None = None) -> tuple[int, str]:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            cwd=cwd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
        out, _ = await proc.communicate()
        return proc.returncode or 0, out.decode(errors="replace")


_manager: McpManager | None = None


def get_manager() -> McpManager:
    global _manager
    if _manager is None:
        _manager = McpManager()
    return _manager